evaluate_modify_with_method.py

# For loading models and data
from torch.utils.data import DataLoader
from src.utils.load_models import load_model
from src.data.datasets import DatasetWithIndices

# Remaining libraries
import torch
import json
import argparse
import importlib
import yaml
import os
import numpy as np

def _apply_method_override(config: dict, method: str, step_size: float | None =None):
    """
    Override the evaluation time-stepping method from CLI without touching the YAML file on disk.
    Policy:
      - dopri5: remove/ignore step_size (set vel_field_method_options = {})
      - euler/rk4: keep only step_size if present
    """
    if method is None:
        return config

    config.setdefault("model", {})
    config["model"].setdefault("RONOM_args", {})

    config["model"]["RONOM_args"]["vel_field_method"] = method

    if method == "dopri5":
        config["model"]["RONOM_args"]["vel_field_method_options"] = {}
    else:
        opts = config["model"]["RONOM_args"].get("vel_field_method_options", {}) or {}
        if step_size is not None:
            config["model"]["RONOM_args"]["vel_field_method_options"] = {"step_size": step_size}
        elif isinstance(opts, dict) and "step_size" in opts:
            config["model"]["RONOM_args"]["vel_field_method_options"] = {"step_size": opts["step_size"]}
        else:
            config["model"]["RONOM_args"]["vel_field_method_options"] = {}

    return config

def evaluate_on_test_or_train(model, dataloader, eval_exp_name, device, train_or_test="test"):
    
    ### Check the the train_or_test input
    assert train_or_test in ["train", "test"], "train_or_test is not 'train' or 'test"
    
    ### Get all reconstruction metrics
    recon_metric_list = []

    ### Save all the reconstructions, ground truths, latent codes, and indices in one list
    recon_list = []
    gt_list = []
    idx_list = []
    latent_list = []

    ### For every test sample, do ...
    with torch.no_grad():
        for u_enc, x_enc, t_ins, u_gt, x_dec, idx in dataloader:
            
            ### Put everything on the correct device
            u_enc, x_enc, t_ins, u_gt, x_dec = u_enc.to(device), x_enc.to(device), t_ins.to(device), u_gt.to(device), x_dec.to(device)

            ### Get the output
            u, latents = model(u_enc=u_enc, x_enc=x_enc, t_ins=t_ins, x_dec=x_dec)
            
            
            ### Save the latents in case the latents are not None (in other words, when there are actually latent codes)
            if latents is not None:
                latent_list.append(latents.detach().cpu().numpy())

            ### Calculate the reconstruction metrics and add it to the list
            if not (u.shape == u_gt.shape):
                print("The shape of the reconstruction {} is not equal to the shape of the ground truth {}. Cannot calculate the reconstruction metric\n and hence we skip the " 
                      "specific evaluation experiment with name {}...".format(u.shape, u_gt.shape, eval_exp_name))
                return None
            else:
                recon_metrics = torch.mean(((u - u_gt) ** 2).flatten(start_dim=1), dim=-1).detach().cpu()
                recon_metrics = torch.sqrt(recon_metrics)
                recon_metric_list.append(recon_metrics)

            ### Save the reconstructions and indices in a list
            recon_list.append(u.detach().cpu().numpy())
            gt_list.append(u_gt.detach().cpu().numpy())
            idx_list.append(idx.detach().cpu().numpy())
    
    ### Get all the reconstruction metrics in one big tensor
    recon_metric = torch.concatenate(recon_metric_list, dim=0)

    ### Calculate relevant quantities
    mean, std = torch.std_mean(recon_metric)
    max = torch.max(recon_metric)
    min = torch.min(recon_metric)
    median = torch.median(recon_metric)

    ### Do the same for the reconstructions (recon_list) and indices (idx_list) AND make them numpy arrays 
    indices = np.concatenate(idx_list)
    
    return mean, std, max, min, median, latent_list, indices
  
  
def evaluate_performance(model, data_class, dataset_args, batch_size_test, 
                         experiment_directory, eval_exp_name, device='cuda'):
    
    ### Make an announncement that we are currently dealing with evaluation eval_exp_name
    num_chars = len(eval_exp_name)
    print("\n")
    print("#############################" + "#"*num_chars + "####")
    print("### Starting with evaluating {} ###".format(eval_exp_name))
    print("#############################" + "#"*num_chars + "####")
    print("\n") 
    ### Create the datasets 
    _, _, _, dataset_test = data_class.generate_datasets(**dataset_args)
    
    ### Modify the datasets
    dataset_test = DatasetWithIndices(dataset_test)

    ### Get a dataloader for the test dataset
    test_dataloader = DataLoader(dataset_test, batch_size=batch_size_test, shuffle=False)

    ### Get the evaluation on the test dataset
    print("\n## Getting reconstruction metrics ## \n")

    out_test = evaluate_on_test_or_train(model, test_dataloader, eval_exp_name, device, train_or_test="test")
    if out_test is None:
        return None
    else:
        ### Get the reconstruction metrics
        mean, std, max, min, median, _, _ = out_test 
        
    ### Finally, return the earlier reconstruction metric quantities
    return mean, std, median, max, min  

def check_test_args(dataset_args, task_name): 
    length_list = -1
    for val in dataset_args.values():
       if not isinstance(val, list):
            raise ValueError("One of the options for testing {} is not a list...".format(task_name))
       elif length_list == -1:
            length_list = len(val)
       elif not len(val) == length_list:
            raise ValueError("One of the options for testing {} is a list that has unequal size to the other lists...".format(task_name))
    return length_list

def evaluate(experiment_directory, method=None, step_size=None):

    ### Get the CPU or GPU on which we will put everything
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    ### Load the YAML file with the training and model specifications
    with open(os.path.join(experiment_directory, "config.yaml"), "r") as file:
        config = yaml.safe_load(file)

    # Override evaluation integrator from CLI (does NOT modify YAML on disk)
    config = _apply_method_override(config, method, step_size)
    
    ### Get the parameters for the dataset
    dataset_type = config["dataset"]["dataset_type"]
    dataset_generator_options = config["dataset"]["dataset_generator_options"]
    dataset_args_train = config["dataset"]["dataset_args_train"]

    ### Get the batch size for testing
    batch_size_test = config["train_config"]["batch_size_test"]
    batch_size_train = batch_size_test 

    ### Get the dataset class
    data_class = getattr(importlib.import_module("src.data.datasets"), dataset_type)(**dataset_generator_options)
    
    ### Load the model
    model = load_model(config, experiment_directory).to(device)
    model.eval()

    ### Print the number of parameters in the model
    print('parameters---------------')
    param_list = []
    for name, param in model.named_parameters():
        if param.requires_grad:
            print(name, param.size())
            param_list.append(param)
    print('-------------------------')
    print('Total para:',sum(p.numel() for p in param_list))
    
    ### Do the evaluation

    ### First, we define a dictionary in which we save all the results. The keys will be the tasks. The values will depend on the task
    recon_metric_dict = {}

    ### Subsequently, we evaluate the performance when evaluating on the training input resolution and map to the output training resolution.
    mean, std, median, max, min = evaluate_performance(model, data_class, dataset_args_train, batch_size_test, 
                                                       experiment_directory=experiment_directory,
                                                       eval_exp_name="input generalization",
                                                       device=device
                                                       )
    recon_metric_dict["input generalization"] = {}
    for val, val_name in zip([mean, std, median, max, min], ["mean", "std", "median", "max", "min"]):
         recon_metric_dict["input generalization"][val_name] = val.item()

    ### Now we evaluate discretization robustness and superresolution. 
    recon_metric_dict["discretization properties"] = {}
    dataset_args_test = config["dataset"]["dataset_args_test"]
    num_discretizations = check_test_args(dataset_args_test, "discretization properties")
    for i in range(num_discretizations):
        dataset_args = {key: value[i] for key, value in dataset_args_test.items()}
        test_task_name = '___'.join([key + "_" + str(value) for key, value in dataset_args.items()])
        recon_metric_dict["discretization properties"][test_task_name] = {}
        performance_metrics = evaluate_performance(model, data_class, dataset_args, batch_size_test, 
                                                           experiment_directory=experiment_directory, 
                                                           eval_exp_name=os.path.join("discretization properties", test_task_name),
                                                           device=device)
        if performance_metrics is None:
            continue
        else:
            mean, std, median, max, min = performance_metrics
        
        for val, val_name in zip([mean, std, median, max, min], ["mean", "std", "median", "max", "min"]):
            recon_metric_dict["discretization properties"][test_task_name][val_name] = val.item()

    ### Now we save the results of the dictionary in a file
    with open(os.path.join(experiment_directory, f"evaluation_results_{method}_{step_size}.json"), "w") as outfile: 
        json.dump(recon_metric_dict, outfile, indent=4, sort_keys=True)

if __name__ == "__main__":
     
    arg_parser = argparse.ArgumentParser(description="Evaluate effects of different ODE solver on RONOM.")
    arg_parser.add_argument(
        "--experiment",
        "-e",
        dest="experiment_directory",
        required=True,
        help="The experiment directory. This directory should include "
        + "experiment specifications in 'config.yaml', and logging will be "
        + "done in this directory as well.",
    )
    arg_parser.add_argument(
        "--method",
        dest="method",
        default=None,
        choices=["euler", "rk4", "dopri5"],
        help="Evaluation integrator override (does not change YAML).",
    )
    arg_parser.add_argument(
        "--step_size",
        dest="step_size",
        type=float,
        default=None,
        help="Step size for evaluation integrators that need it (euler, rk4). Ignored if method is dopri5 or None.",
    )

    args = arg_parser.parse_args()

    evaluate(args.experiment_directory, method=args.method, step_size=args.step_size)