evaluate_computational_costs.py

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

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

# Libraries needed for dealing with the DeepONet
import deepxde as dde
import math


### Stuff for loading the DeepONet
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 _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_computational_costs(experiment_directory, method=None, step_size=None, subsample_space=None, subsample_time=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). Also only do this when you deal with RONOM
    if "model_type" in config and config["model_type"] == "RONOM":
        config = _apply_method_override(config, method, step_size)

    ### Override the subsampling options
    if subsample_space is not None:
        if "model_type" in config:
            config["dataset"]["dataset_args_train"]["subsampling_x_input"] = subsample_space
        config["dataset"]["dataset_args_train"]["subsampling_x_output"] = subsample_space
    if subsample_time is not None:
        config["dataset"]["dataset_args_train"]["subsampling_t"] = subsample_time

    ### 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 dataset class
    data_class = getattr(importlib.import_module("src.data.datasets"), dataset_type)(**dataset_generator_options)
    
    ### Create the datasets 
    dataset_sup_train, dataset_unsup_train, _, dataset_test = data_class.generate_datasets(**dataset_args_train)
    
    ### Modify the datasets
    dataset_test = DatasetWithIndices(dataset_test)

    ### Load the model
    if "model_type" in config:
        model = load_model(config, experiment_directory).to(device)
        model.eval()
    else: # In this case, we are dealing with a DeepONet
        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()

    ### Initialize a list to store the computational times
    computational_times = []
    
    ### Grab a single test case and use that throughout
    u_enc, x_enc, t_ins, u_gt, x_dec, _ = dataset_test[0]
    u_enc = u_enc[None, :]
    x_enc = x_enc[None, :]
    t_ins = t_ins[None, :]
    u_gt = u_gt[None, :]
    x_dec = x_dec[None, :]
    
    num_repititions = 230
    
    ### If you have a neural operator or RONOM
    if "model_type" in config:
  
        starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
        computational_times = []
        
        #GPU-WARM-UP
        for _ in range(100):
            _, _ = model(u_enc=u_enc, x_enc=x_enc, t_ins=t_ins, x_dec=x_dec)
        
        with torch.no_grad():
            for i in range(num_repititions):
                starter.record()
                ### Get the output
                _, _ = model(u_enc=u_enc, x_enc=x_enc, t_ins=t_ins, x_dec=x_dec)
                ender.record()
                
                # WAIT FOR GPU SYNC
                torch.cuda.synchronize()
                curr_time = starter.elapsed_time(ender)
                computational_times.append(curr_time/1000)
    
    else:

        starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
        computational_times = []
        
        #GPU-WARM-UP
        for _ in range(100):
            
            # 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
            _ = model([u_enc[..., 0].flatten(start_dim=1), xt[0, ...].flatten(end_dim=-2)]).reshape(u_enc.shape[0], *xt_shape[1:-1])[..., None]       

        ### Do the same thing num_repititions times
        with torch.no_grad():
            for i in range(num_repititions):
                
                starter.record()
                
                # 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
                _ = model([u_enc[..., 0].flatten(start_dim=1), xt[0, ...].flatten(end_dim=-2)]).reshape(u_enc.shape[0], *xt_shape[1:-1])[..., None]

                ender.record()
                
                # WAIT FOR GPU SYNC
                torch.cuda.synchronize()
                curr_time = starter.elapsed_time(ender)
                computational_times.append(curr_time/1000)
    
    ## Define the filename
    extension=""
    if "model_type" in config and config["model_type"] == "RONOM":
        extension = "_" + config["model"]["RONOM_args"]["vel_field_method"]
        extension = extension + "_{}".format(step_size)
    extension = extension + "_xsub{}_tsub{}".format(subsample_space, subsample_time)
    filename = "computational_costs"+extension+".txt"
    
    ### Save the the mean, median, standard deviation, maximum, and minimum of the computational times to a text file
    computational_times_np = np.array(computational_times)
    mean_time = np.mean(computational_times_np)
    median_time = np.median(computational_times_np)
    std_time = np.std(computational_times_np)
    max_time = np.max(computational_times_np)
    min_time = np.min(computational_times_np)
    with open(os.path.join(experiment_directory, filename), "w") as f:
        f.write(f"Mean time: {mean_time:.6f} seconds\n")
        f.write(f"Median time: {median_time:.6f} seconds\n")
        f.write(f"Standard deviation: {std_time:.6f} seconds\n")
        f.write(f"Maximum time: {max_time:.6f} seconds\n")
        f.write(f"Minimum time: {min_time:.6f} seconds\n")
    

if __name__ == "__main__":
     
    arg_parser = argparse.ArgumentParser(description="Evaluate the computational costs of RONOM, FNO, CNO, or 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.",
    )
    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.",
    )
    arg_parser.add_argument(
        "--subsample_space",
        dest="subsample_space",
        type=int,
        default=None,
        help="The subsampling done in space for the input and output."
    )
    arg_parser.add_argument(
        "--subsample_time",
        dest="subsample_time",
        type=int,
        default=None,
        help="The subsampling done in time for the output."
    )
    args = arg_parser.parse_args()

    evaluate_computational_costs(args.experiment_directory, method=args.method, step_size=args.step_size, 
                                 subsample_space=args.subsample_space, subsample_time=args.subsample_time)