evaluate_deeponet.py

# For loading models and data
from torch.utils.data import DataLoader
from src.data.datasets import DatasetWithIndices, ConcatenatedDataset
import deepxde as dde

# Get the functions required for plotting
from src.visualization.plotting_solutions import plot_2d_recon_vs_gt, plot_2d
from src.visualization.plotting_latent_space import visualize_latent_space

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

def random_fourier_encoding(xt, B):
    x, t = xt[..., :-1], xt[..., -1]
    x = 2 * math.pi * (x @ B.T)
    return torch.concatenate((torch.cos(x), torch.sin(x), t[..., None]), axis=-1)

def load_deeponet_model(experiment_directory, input_size_branch_net, pde_type, model_config, device, B=None):
    
    features_trunk = model_config["features_trunk"]
    features_branch = model_config["features_branch"]
    num_basis_funcs = model_config["num_basis_funcs"]
    activation = model_config["activation"]
    weight_initializer = model_config["weight_initializer"]
    
    if pde_type == "Burgers" or pde_type == "WaveEq" or pde_type == "NavierStokes":
        
        spatial_dim = 1 if pde_type == "Burgers" else 2

        # Define the model
        if model_config["positional_encoding"]["is_used"]: 
            positional_encoding_size = model_config["positional_encoding"]["dim"]
            net = dde.nn.DeepONetCartesianProd(
                [input_size_branch_net, *features_branch, num_basis_funcs],
                [2 * positional_encoding_size+1, *features_trunk, num_basis_funcs],
                activation,
                weight_initializer
            )
            if B is None:
                B = torch.normal(0, 1, size=(positional_encoding_size, spatial_dim)).to(device)
            net.apply_feature_transform(lambda xt: random_fourier_encoding(xt, B))    
        else:
            net = dde.nn.DeepONetCartesianProd(
                [input_size_branch_net, *features_branch, num_basis_funcs],
                [2, *features_trunk, num_basis_funcs],
                activation,
                weight_initializer
            )
        model = net
        
        # Load the model weights into the model
        weights = torch.load(os.path.join(experiment_directory, "ModelParameters", "model_params"), weights_only=True)["model_state_dict"]
        model.load_state_dict(weights)
        
        # Put the model on the correct device and return everything
        model.to(device)
        return model
    
    else:
        raise NotImplementedError("Currently, we only implemented the Burgers, Wave, and Navier-Stokes equation...")
    
def evaluate_on_test_or_train(model, dataloader, experiment_directory, 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)
            
            # Concatenate x_dec and t_ins
            xt = torch.cat([x_dec[:, None, ...].repeat(1, t_ins.shape[-1], *((x_dec.ndim-1)*(1,))), t_ins[:, :, *((x_dec.ndim-1)*(None,))].repeat(1, 1, *x_dec.shape[1:-1], 1)], dim=-1)
            
            # Get the prediction of the model (assuming we evaluate the prediction always on the same samples!)
            xt_shape = xt.shape
            u = model([u_enc[..., 0].flatten(start_dim=1), xt[0, ...].flatten(end_dim=-2)]).reshape(u_enc.shape[0], *xt_shape[1:-1])[..., None]

            ### Put the latent codes to None. The reason we do this is because the rest of the code is just as in evaluate.py which requires latent codes to be defined.
            latents = None
            
            ### 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
            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 
    recons = np.concatenate(recon_list)
    gts = np.concatenate(gt_list)
    indices = np.concatenate(idx_list)

    ### Plot the reconstructions and save the figures to the correct folder
    save_dir = os.path.join(experiment_directory, "figures", eval_exp_name, train_or_test)
    if not os.path.isdir(save_dir):
        os.makedirs(save_dir)
    for i in [*list(range(0, recons.shape[0], 80)), torch.argmax(recon_metric).item()]:  # Plot every 40th sample to avoid too many plots
        
        # Get the name of the sample
        if i == torch.argmax(recon_metric).item():
            sample_name = "Worst reconstruction"
        else:
            sample_name = indices[i]
    
        if len(recons[i, ...].squeeze().shape) == 2:
            
            # Save the comparison between the reconstruction and the ground truth
            plot_2d_recon_vs_gt(recons[i, ...].squeeze(), gts[i, ...].squeeze(),
                                boundaries_x=[0.0, 1.0], boundaries_y=[0.0, 1.0],
                                save_dir=save_dir, save_name="Test sample {}".format(sample_name), 
                                x_axis_name="t", y_axis_name="x", title="u(x,t)")
            
            # Finally, also save the ground truth and reconstructed images
            np.save(os.path.join(save_dir, "Test sample {}.npy".format(sample_name)), recons[i, ...].squeeze())
            np.save(os.path.join(save_dir, "ground truth {}.npy".format(sample_name)), gts[i, ...].squeeze())
            
            # Also save and plot the error
            plot_2d(gts[i, ...].squeeze() - recons[i, ...].squeeze(), boundaries_x=[0.0, 1.0], boundaries_y=[0.0, 1.0],
                save_dir=save_dir, save_name="Error sample {}".format(sample_name), x_axis_name="t", y_axis_name="x", title="u(x,t)")
            
        elif len(recons[i, ...].squeeze().shape) == 3:
            
            for j in range(recons[i, ...].shape[0]):
                
                # Get the main directory for the test sample
                dirname = os.path.join(save_dir, "Test sample {}".format(sample_name))
                if not os.path.isdir(dirname):
                    os.makedirs(dirname)
                    
                # Plot the ground truth and reconstruction for each timepoint side-by-side
                comparison_save_dir = os.path.join(dirname, "comparing_gt_vs_reconstruction")
                if not os.path.isdir(comparison_save_dir):
                    os.makedirs(comparison_save_dir)
                plot_2d_recon_vs_gt(recons[i, j, ...].squeeze(), gts[i, j, ...].squeeze(),
                                    boundaries_x=[0.0, 1.0], boundaries_y=[0.0, 1.0],
                                    save_dir=comparison_save_dir, save_name="timepoint {}".format(j), 
                                    x_axis_name="x", y_axis_name="y", title="u([x,y],t)")
                
                # Also plot them individually for the reconstruction
                recon_save_dir = os.path.join(dirname, "reconstruction")
                if not os.path.isdir(recon_save_dir):
                    os.makedirs(recon_save_dir)
                plot_2d(recons[i, j, ...].squeeze(), boundaries_x=[0.0, 1.0], boundaries_y=[0.0, 1.0],
                save_dir=recon_save_dir, save_name="reconstruction timepoint {}".format(j), hide_axes_ticks=True,
                max_val=None, min_val=None)
                
                # Also do this for the ground truth
                gt_save_dir = os.path.join(dirname, "ground truth")
                if not os.path.isdir(gt_save_dir):
                    os.makedirs(gt_save_dir)
                plot_2d(gts[i, j, ...].squeeze(), boundaries_x=[0.0, 1.0], boundaries_y=[0.0, 1.0],
                save_dir=gt_save_dir, save_name="ground truth timepoint {}".format(j), hide_axes_ticks=True,
                max_val=None, min_val=None)
                
                # Also do this for the error
                bot = gts[i, j, ...].squeeze() - recons[i, j, ...].squeeze()
                bot_save_dir = os.path.join(dirname, "error")
                if not os.path.isdir(bot_save_dir):
                    os.makedirs(bot_save_dir)
                plot_2d(bot, boundaries_x=[0.0, 1.0], boundaries_y=[0.0, 1.0], 
                        save_dir=bot_save_dir, save_name="error timepoint {}".format(j), hide_axes_ticks=True,
                        max_val=None, min_val=None)
                
                # Finally, also save the ground truth and reconstructed images
                np.save(os.path.join(dirname, "reconstruction", "reconstruction timepoint {}.npy".format(j)), recons[i, j, ...].squeeze())
                np.save(os.path.join(dirname, "ground truth", "ground truth timepoint {}.npy".format(j)), gts[i, j, ...].squeeze())
                
            # else:
            #     break # Only plot the reconstructions of the first 5 samples
        else:
            raise ValueError("The number of dimensions of the reconstruction is not 2 or 3, but {}. Cannot plot the reconstruction...".format(len(recons[i, ...].squeeze().shape)))

    return mean, std, max, min, median, latent_list, indices
  
  
def evaluate_performance(model, data_class, dataset_args, batch_size_train, batch_size_test, 
                         experiment_directory, eval_exp_name, calculate_lat_embedding=True, 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_sup_train, dataset_unsup_train, _, dataset_test = data_class.generate_datasets(**dataset_args)
    
    ### Modify the datasets
    dataset_train = DatasetWithIndices(ConcatenatedDataset(dataset_sup_train, dataset_unsup_train))
    dataset_test = DatasetWithIndices(dataset_test)

    ### Get a dataloader for the test dataset
    test_dataloader = DataLoader(dataset_test, batch_size=batch_size_test, shuffle=False, generator=torch.Generator(device=device))

    ### Get the evaluation on the test dataset
    print("\n## Getting reconstruction metrics ## \n")
    train_dataloader = DataLoader(dataset_train, batch_size=batch_size_train, shuffle=False, generator=torch.Generator(device=device))
    #mean, std, max, min, median, latents_test_list, indices_test = evaluate_on_test_or_train(model, train_dataloader, experiment_directory, eval_exp_name, device, train_or_test="train")
    
    mean, std, max, min, median, latents_test_list, indices_test = evaluate_on_test_or_train(model, test_dataloader, experiment_directory, eval_exp_name, device, train_or_test="test")
    
    ### In case the list of latent codes is not empty (aka, we are dealing with a latent model), get some plots of the latent space
    if not len(latents_test_list) == 0 and calculate_lat_embedding:
        
        #############################################################################################
        ### This code is currently not used and may be outdated. It has not been tested recently. ###
        ### It is kept here for reference, as it may be useful in the future.                     ###
        ### The DeepONet does not explicitly use latent codes, so latent code visualization is    ###
        ### not even applicable here. However, it is kept for consistency with `evaluate.py`.     ###
        #############################################################################################
        
        print("## Getting the latent plots ##\n")
        
        ### Define the train dataloader
        train_dataloader = DataLoader(dataset_train, batch_size=batch_size_train, shuffle=False, generator=torch.Generator(device=device))
        
        ### Get the latent codes
        _, _, _, _, _, latents_train_list, indices_train = evaluate_on_test_or_train(model, train_dataloader, experiment_directory, eval_exp_name, device, train_or_test="train")
        
        ### Create the latent space plots
        latents_train = np.concatenate(latents_train_list)
        latents_test = np.concatenate(latents_test_list)
        visualization_types = ["PCA", "UMAP", "PHATE"]
        time_based_colorings = [True, False]
        plot_types = ["train", "test", "joint"]
        for visualization_type in visualization_types:
            visualize_latent_space(latents_train, latents_test, indices_train, indices_test, visualization_type, 
                                exp_dir=experiment_directory, time_based_colorings=time_based_colorings, plot_types=plot_types)
        
    ### 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):

    ### 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)

    ### 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
    dataset_sup_train, dataset_unsup_train, _, _ = data_class.generate_datasets(**dataset_args_train)
    input_size_branch_net = math.prod((dataset_sup_train.__getitem__(0)[0] if len(dataset_sup_train) > 0 else dataset_unsup_train.__getitem__(0)[0]).shape[:-1])
    B = torch.load(os.path.join(experiment_directory, "ModelParameters", "B.pt"), weights_only=True).to(device) if os.path.isfile(os.path.join(experiment_directory, "ModelParameters", "B.pt")) else None
    model = load_deeponet_model(experiment_directory, input_size_branch_net, dataset_type, config["model"], device, B)
    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_train, batch_size_test, 
                                                       experiment_directory=experiment_directory,
                                                       eval_exp_name="input generalization",
                                                       calculate_lat_embedding=True,
                                                       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()])
        if not (dataset_args["subsampling_x_input"] == dataset_args_train["subsampling_x_input"]):
            print("Skipping task {} because the subsampling_x_input is not equal to the one used for training...".format(test_task_name))
            continue
        recon_metric_dict["discretization properties"][test_task_name] = {}
        mean, std, median, max, min = evaluate_performance(model, data_class, dataset_args, batch_size_train, batch_size_test, 
                                                           experiment_directory=experiment_directory, 
                                                           eval_exp_name=os.path.join("discretization properties", test_task_name),
                                                           calculate_lat_embedding=False,
                                                           device=device)
        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, "evaluation_results.json"), "w") as outfile: 
        json.dump(recon_metric_dict, outfile, indent=4, sort_keys=True)

if __name__ == "__main__":
     
    arg_parser = argparse.ArgumentParser(description="Evaluate DeepOnet on a given dataset.")
    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.",
    )
    args = arg_parser.parse_args()

    evaluate(args.experiment_directory)