planning_dia.py

# ### Initializing
from utils.utils import *
from dataset_generator import run
import pandas as pd
import json, argparse, warnings, yaml
import torch
# torch.set_default_dtype(torch.float64) 
# torch.backends.cudnn.allow_tf32 = False
# torch.backends.cuda.matmul.allow_tf32 = False
warnings.filterwarnings('ignore')
from train_test import *
from models.models import *
import pickle, os, time




def evaluate(wds_test, model, reservoirs, args, save_dir, out_f, e=1e-32, llist=False):
    """ Initializing hyperparameters. """
    n_epochs, learn_r, alpha = args.n_epochs, args.lr, args.weight_decay


    """ Loading the trained model. """
    if args.model_path is None:
        args.model_path = save_dir+"/model_"+args.model+"_"+str(args.n_epochs)+"_"+str(args.n_aggr)+\
                            "_"+str(args.rand_sensors)+".pt"
    model_state = torch.load(args.model_path)
    model.load_state_dict(model_state["model"])

    # model = model.double()
    
    """ Initializing parameters and loading dataset and batches. """
    n_nodes = wds_test.X.shape[1]
    n_edges = wds_test.edge_attr[0].shape[0]


    test_dataset, Y_test = load_dataset(wds_test, n_nodes, reservoirs)
    test_loader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False)


    """" Testing and saving the results. """
    batch = next(iter(test_loader))

    batch.edge_attr = batch.edge_attr.to(device)
    diameter = batch.edge_attr[:n_edges//2, 1:2].clone()
    diameter_star = diameter.clone()
    diameter = torch.nn.Parameter(data=diameter).to(device)

    
    args.h_max = args.h_max.repeat(args.batch_size, 1).to(device)
    args.h_min = args.h_min.repeat(args.batch_size, 1).to(device)
    delta_h = (args.h_max - args.h_min)
    
    optimizer = Adam([diameter], lr=learn_r, weight_decay=alpha)    
    lr_decay_step, lr_decay_rate = 10, .75
    opt_scheduler = lr_scheduler.MultiStepLR(optimizer, range(lr_decay_step, lr_decay_step*1000, lr_decay_step), gamma=lr_decay_rate)
         
    
    model.eval()
    
    for epoch in tqdm(range(n_epochs)):
    
        test_losses, loss_dias, loss_constraintss, model_losses = [], [], [], []
        Y_hat, F_hat, D_hat, D = [], [], [], []

        for batch in test_loader:

            optimizer.zero_grad()

            batch.edge_attr = batch.edge_attr.to(device)

            batch.edge_attr = torch.stack(batch.edge_attr.split(n_edges))
            r_in = (10.667 * batch.edge_attr[:, :, 0:1]) * \
                    (torch.pow(batch.edge_attr[:, :, 2:3], -1.852) * \
                    torch.pow(diameter.relu(), -4.871).repeat(2, 1))
            r_in = torch.nan_to_num(r_in, nan=0, posinf=0, neginf=0).reshape(-1, 1)
        
           
            y_hat = model(batch, r_iter=args.r_iter, zeta=e, r_in=r_in)

            loss_dia = (diameter_star - diameter).abs().mean()

            loss_constraints = torch.maximum((y_hat[~model.reservoir_mask] - args.h_max[~model.reservoir_mask]), torch.zeros_like(y_hat[~model.reservoir_mask])).mean() + \
                                torch.maximum((args.h_min[~model.reservoir_mask] - y_hat[~model.reservoir_mask]), torch.zeros_like(y_hat[~model.reservoir_mask])).mean()

            _loss = args.alpha * model.loss() + args.beta * loss_dia + args.gamma * loss_constraints 

            _loss.backward()

            grad_norm = torch.nn.utils.clip_grad_norm_([diameter], args.clip_val)
            
            pre_update_diameter = diameter.data.clone()

            optimizer.step()

            mask = diameter_star > diameter
            diameter.data[mask] = pre_update_diameter[mask]
            
            test_losses.append(_loss.detach().cpu().item())
            loss_dias.append((diameter_star - diameter).abs().sum().detach().cpu().item())
            loss_constraintss.append(loss_constraints.detach().cpu().item())
            model_losses.append(model.loss().detach().cpu().item())

            D.append(model.d_star.detach().cpu())

            Y_hat.append(torch.hstack(y_hat.detach().cpu().split(n_nodes)).view(n_nodes, -1, y_hat.shape[1]))

            F_hat.append(model.q_hat.detach().cpu())

            D_hat.append(torch.hstack(model.d_hat.detach().cpu().split(n_nodes)).view(n_nodes, -1, model.d_hat.shape[1]))

        opt_scheduler.step()        

        Y_hat = torch.cat(Y_hat, dim=1).transpose(1,0)
        D_hat = torch.cat(D_hat, dim=1).transpose(1,0).abs()

        res_mask = torch.zeros_like(Y_hat)
        res_mask[:, reservoirs, :] = Y_hat[:, reservoirs, :]
        res_mask = res_mask.mean(0).mean(1).bool()

        print("Epoch: ", epoch, "Test loss: ", np.round(np.mean(test_losses), 8), 
                "\t loss_dia: ", np.round(np.mean(loss_dias), 8), 
                "\t loss_constraints: ", np.round(np.mean(loss_constraintss), 8), 
                "\t loss_model: ", np.round(np.mean(model_losses), 8), 
                '\t Grad Norm: ', grad_norm.mean().cpu().numpy().round(4).item())
        print("Epoch: ", epoch, "Test loss: ", np.round(np.mean(test_losses), 8), 
                "\t loss_dia: ", np.round(np.mean(loss_dias), 8), 
                "\t loss_constraints: ", np.round(np.mean(loss_constraintss), 8), 
                "\t loss_model: ", np.round(np.mean(model_losses), 8), 
                '\t Grad Norm: ', grad_norm.mean().cpu().numpy().round(4).item(), 
                file=out_f)
        print("max head: ", Y_hat.max().item(), "min_head: ", Y_hat.min().item(), 
                "min_pressure: ", (Y_hat[:, ~res_mask, :] - args.elevs[~res_mask]).min().item(), 
                "mean_pressure: ", (Y_hat[:, ~res_mask, :] - args.elevs[~res_mask]).mean().item())
        print("max head: ", Y_hat.max().item(), "min_head: ", Y_hat.min().item(), 
                "min_pressure: ", (Y_hat[:, ~res_mask, :] - args.elevs[~res_mask]).min().item(), 
                "mean_pressure: ", (Y_hat[:, ~res_mask, :] - args.elevs[~res_mask]).mean().item(),
                file=out_f)




    print("Test loss: ", np.round(np.mean(test_losses), 8))
    print("Test loss: ", np.round(np.mean(test_losses), 8), file=out_f)
        
    return Y_test, Y_hat, test_losses, F_hat, D_hat, D, diameter.detach().cpu()



if __name__ == '__main__':



    """ Creating directories. """
    file_dir = os.getcwd()
    if not os.path.isdir(os.path.join(file_dir, "trained_models_2024")):
        os.system('mkdir ' + os.path.join(file_dir, "trained_models_2024"))
    save_dir = os.path.join(file_dir, "trained_models_2024")
    if not os.path.isdir(save_dir):
        os.system('mkdir ' + save_dir)

    """ Creating an output file to log progress. """
    out_f = open(save_dir+"/output_"+str(datetime.date.today())+".txt", "a")

    with open('args.json', 'r') as args_file:
        args_dict = json.load(args_file)
    args = argparse.Namespace(**args_dict)

    args.wdn = "hanoi"
    args.div_factor = 1

    s = 2000



    # ### Reading Scenario


    X_tvt, edge_indices_tvt, edge_attr_tvt, wntr_demands = [], [], [], []

    scenario_path = os.path.join(os.getcwd(),"networks",  args.wdn, "real", "s"+str(s))
    args.inp_file = os.path.join(scenario_path, args.wdn + ".inp")
    
    args.path_to_data = os.path.join(scenario_path, "Results-Clean", "Measurements_All.xlsx")

    wdn_graph, reservoirs = create_graph(args.inp_file, args.path_to_data, ldc=True)

    wntr_d_df = pd.read_csv(os.path.join(scenario_path, "Results-Clean", "Measurements_All_Demands.csv"))
    wntr_d_df['Timestamp'] = pd.to_datetime(wntr_d_df['Timestamp'])#, unit='s')
    wntr_d_df = wntr_d_df.set_index("Timestamp")
    wntr_d = torch.tensor(wntr_d_df.astype("float32").values)[::args.div_factor]
    wntr_demands.append(wntr_d)

    wn = wntr.network.WaterNetworkModel(args.inp_file)   
    base_demand = torch.zeros((wdn_graph.X.shape[1])).float() 
    res_mask = torch.zeros((wdn_graph.X.shape[1])).bool()
    res_mask[reservoirs] = 1
    base_demand[~res_mask] = torch.tensor(wn.query_node_attribute('base_demand')).float()
    base_demand = base_demand.unsqueeze(0).unsqueeze(2).repeat(wdn_graph.X.shape[0], 1, 1)
    wdn_graph.X = torch.cat((wdn_graph.X, base_demand), dim=-1)

    """ Creating train-val-test data based on the specified number of samples. """
    X_s = wdn_graph.X.clone()
    edge_indices_s = wdn_graph.edge_indices.clone()
    edge_attr_s = wdn_graph.edge_attr.clone()

    """ Creating train-val-test splits. """
    X_tvt.append(X_s)
    edge_indices_tvt += list(edge_indices_s)
    edge_attr_tvt += list(edge_attr_s)


    print('\t', str(s))

    X_tvt = torch.vstack(X_tvt)    

    wds_tvt = WDN_Graph(X=X_tvt, edge_indices=edge_indices_tvt, edge_attr=edge_attr_tvt)


    n_nodes = wds_tvt.X.shape[1]
    n_edges = wds_tvt.edge_indices[0].shape[1]
    n_samples = wds_tvt.X.shape[0]

    print(wds_tvt.X.shape, wds_tvt.edge_indices[0].shape, wds_tvt.edge_attr[0].shape)



    args.model_path = "trained_models_2024/" + args.wdn + "/model_" + args.wdn + ".pt"
    args.r_iter = 20
    args.n_mlp = 2

    args.model == "PI_GCN"

    model = PI_GCN( M_n = 2,                    # number of node features (d_star, d_hat).
                    out_dim = 1,                # out dimension is 1 since only flows are directly estimated.
                    M_e = 2,                    # number of edge features (q_hat, q_tilde).
                    M_l = args.M_l,             # specified latent dimension.
                    I = 5,                      # number of GCN layers.
                    num_layers = args.n_mlp,    # number of NN layers used in every MLP.
                    n_iter = args.n_iter,       # minimum number of iterations.
                    bias = False                # we do not use any bias.
                    ).to(device)


    model_state = torch.load(args.model_path)
    model.load_state_dict(model_state["model"])
    e = 1e-12

    wn = wntr.network.WaterNetworkModel(args.inp_file)
    elevs_junctions = torch.tensor(wn.query_node_attribute('elevation')).float()[:, None]
    args.elevs = torch.zeros((n_nodes, 1))
    args.elevs[:elevs_junctions.shape[0], :] = elevs_junctions


    wds_eval = copy.deepcopy(wds_tvt)


    eval_samples = wds_eval.X[..., 1].max(dim=0)[1].unique()
    print(eval_samples)



    wds_eval.X = wds_eval.X[eval_samples]
    wds_eval.edge_indices = list(torch.stack(wds_eval.edge_indices)[eval_samples])
    wds_eval.edge_attr = list(torch.stack(wds_eval.edge_attr)[eval_samples])

    print(wds_eval.X.shape, len(wds_eval.edge_indices), len(wds_eval.edge_attr))

    args.batch_size = 1 # len(eval_samples)
    args.n_epochs = 100
    args.lr = 1e-3
    args.weight_decay = 0
    args.clip_val = 1e-5

    #                   alpha, beta, gamma                    p_min
    # hanoi              (10., .1, 1.)                         40.
    # fossolo            (100., .1, 1.)                        52.5
    # pescara            (100., 1., 1.)                        20.  
    # l_town             (100., .01, 1.)                      115.
    # zhijiang           (10., 1., 1.)                         30.  
    th_dict = {
                "hanoi" : 40.,
                "fossolo" : 52.5,
                "pescara" : 20.,
                "l_town" : 115.,
                "zhijiang" : 30.,
                }           
    th = th_dict[args.wdn]
    print(th)

    args.alpha = 10.
    args.beta = .1
    args.gamma = 1.

    print(args.alpha, args.beta, args.gamma)

    args.h_max, args.h_min = torch.ones((n_nodes, 1)) * wds_eval.X[:, reservoirs, :].max() , args.elevs + th
    args.h_min[reservoirs] = wds_eval.X[0, reservoirs, 0:1]
    print(args.h_max.min(), args.h_min.min(), args.elevs.max())
    


    start = time.time()

    Y_in, Y_hat, test_losses, F_hat, D_hat, D, diameter  = evaluate(wds_eval, model, reservoirs, args, save_dir, out_f, e=e)
    eval_time = np.round(time.time() - start, 2)
    print('Time taken for ', len(eval_samples), ' samples: ', eval_time, ' seconds')

    dia_change = (diameter - wds_eval.edge_attr[0][:n_edges//2, 1:2]).abs()

    print("Total change in diameters: ", dia_change.sum().item())
    print("Total change in diameters: ", dia_change.sum().item(), file=out_f)
    print("Mean change in diameters: ", dia_change.mean().item())
    print("Mean change in diameters: ", dia_change.mean().item(), file=out_f)

    date_dir = str(datetime.date.today())

    file_dir = os.getcwd() 
    saving_dir = os.path.join(file_dir, "data/"+date_dir)
    if not os.path.isdir(saving_dir):
        os.system('mkdir ' + saving_dir)

    save_dict = {}
    save_dict["Total change in diameters: "] = dia_change.sum().item()
    save_dict["Mean change in diameters: "] = dia_change.mean().item()
    save_df = pd.DataFrame(save_dict, index=[args.wdn]).transpose()
    save_df.to_csv(saving_dir+"/planning_dia_" + args.wdn + "_" + str(s) + "_" + str(len(eval_samples)) + "_dia_dict.csv")

    with open(saving_dir+"/planning_dia_" + args.wdn + "_" + str(s) + "_" + str(len(eval_samples)) + "_wds_eval.pickle", "wb") as _file:
        pickle.dump(wds_eval, _file)

    with open(saving_dir+"/planning_dia_" + args.wdn + "_" + str(s) + "_" + str(len(eval_samples)) + "_Y_in.pickle", "wb") as _file:
        pickle.dump(Y_in, _file)
    with open(saving_dir+"/planning_dia_" + args.wdn + "_" + str(s) + "_" + str(len(eval_samples)) + "_Y_hat.pickle", "wb") as _file:
        pickle.dump(Y_hat, _file)
    with open(saving_dir+"/planning_dia_" + args.wdn + "_" + str(s) + "_" + str(len(eval_samples)) + "_D_hat.pickle", "wb") as _file:
        pickle.dump(D_hat, _file)
    with open(saving_dir+"/planning_dia_" + args.wdn + "_" + str(s) + "_" + str(len(eval_samples)) + "_F_hat.pickle", "wb") as _file:
        pickle.dump(F_hat, _file)
    with open(saving_dir+"/planning_dia_" + args.wdn + "_" + str(s) + "_" + str(len(eval_samples)) + "_diameter.pickle", "wb") as _file:
        pickle.dump(diameter, _file)