planning_dia_gad.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
import pygad
import numpy
from dataset_generator import simulate_wntr




def fitness_func(ga_instance, solution, solution_idx):

    solution_clipped = solution.clip(function_inputs[:,0], None)
    
    output1 = (solution_clipped.reshape(-1,1)) #+ function_inputs
    new_dias = torch.tensor(output1, dtype=torch.float32)
    

    _simulate = True
    demand_series = wds_eval.X[:, :, 1].clone()
    n_eval_samples = len(eval_samples)
    time_delta = datetime.timedelta(minutes = (n_eval_samples) * 30)
    start_time="2018-01-01 00:00"
    end_time = datetime.datetime.strptime(start_time, '%Y-%m-%d %H:%M') + time_delta

    output2 = simulate_wntr(args, 
                            s=s, 
                            s_offset=2002, 
                            demand_series=demand_series, 
                            diameter=new_dias,
                            dem_multiplier=None,  
                            _simulate=_simulate,
                            start_time=start_time, 
                            end_time=end_time,
                            _print=False,
                            )
    
    output2 = output2[:, :, None].astype(np.float32)
    
    print(np.mean(np.abs(output1 - function_inputs)), 
          (output2-args.elevs.numpy()).min(), 
          args.h_min.min().item(), 
          output2.max(), 
          args.h_max.max().item())

    e = 1e-6
    # fitness1 = 1.0 / (np.sum(numpy.abs(output1 - function_inputs)) + e)
    # fitness2 = 1.0 / (np.sum(np.maximum((output2 - args.h_max.numpy()), np.zeros_like(output2)))  + e)
    # fitness3 = 1.0 / (np.sum(np.maximum((args.h_min.numpy() - output2), np.zeros_like(output2)))  + e)
    # fitness4 = 1.0 / (np.sum(np.maximum(0, -1 * output1)) + e)

    fitness1 = -(np.sum(numpy.abs(output1 - function_inputs)))
    fitness2 = np.sum(((output2 - args.h_max.numpy()) > 0) * -1)       
    fitness3 = np.sum(((args.h_min.numpy() - output2) > 0) * -1)       
    fitness4 = np.sum((solution < 0) * -1)       

    print('Solution : ', fitness1, fitness2, fitness3, fitness4)
    # return [fitness1 + fitness2 + fitness3 + fitness4]
    return [fitness1, fitness2, fitness3, fitness4]





""" 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 = "pescara"
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))
# if args.wdn == 'hanoi':
#     args.inp_file = os.path.join(scenario_path, "Hanoi_CMH_Scenario-" + str(s) + ".inp")
# else:
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]#.unsqueeze(0)
wds_eval.edge_indices = list(torch.stack(wds_eval.edge_indices)[eval_samples])#.unsqueeze(0)
wds_eval.edge_attr = list(torch.stack(wds_eval.edge_attr)[eval_samples])#.unsqueeze(0)

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

th_dict = {
                "hanoi" : 40.,
                "fossolo" : 52.5,
                "pescara" : 20.,
                "l_town" : 115.,
                "zhijiang" : 30.,
                }           
th = th_dict[args.wdn]
print(th)

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.shape, args.h_min.shape, args.h_max.min(), args.h_min.min())




function_inputs = wds_eval.edge_attr[0][:n_edges//2, 1:2].numpy()


# pop_range = np.arange(0., .1, .01)
pop_range = np.zeros(100)

initial_population = function_inputs + pop_range

initial_population[:, 5:] += np.random.uniform(0, 1, size=initial_population[:, 5:].shape)
print(initial_population.shape)

num_parents_mating = 50

num_generations = 100

num_genes = function_inputs.shape[0]

dia_list = []
s_time = time.time()
_seed_range = range(21, 22, 11)


#                   random_mutation_max_val                   random_mutation_min_val
# hanoi              0.1                                         -0.01
# fossolo            0.1                                         -0.01
# pescara            0.15                                         -0.01  
# l_town             0.1                                         -0.01
# zhijiang           0.1                                         -0.01  

random_mutation_min_val = -.01
random_mutation_max_val = .15

for _seed in tqdm(_seed_range):
    print(_seed)

    ga_instance = pygad.GA(num_generations = num_generations,
                        initial_population=initial_population.T,
                        random_mutation_min_val = random_mutation_min_val,
                        random_mutation_max_val = random_mutation_max_val,
                        save_best_solutions = True,
                        num_parents_mating = num_parents_mating,
                        num_genes = num_genes,
                        fitness_func = fitness_func,
                        parent_selection_type = 'nsga2',
                        random_seed = _seed,
                        parallel_processing=['process', 12]
    )

    ga_instance.run()

    solution, solution_fitness, solution_idx = ga_instance.best_solution(ga_instance.last_generation_fitness)
    print(f"Fitness value of the best solution = {solution_fitness}")


    new_dias = solution.reshape(-1,1) 

    new_dias[new_dias < function_inputs] = function_inputs[new_dias < function_inputs]

    if np.mean(solution_fitness[1:]) == 0.:
        dia_list.append(new_dias)

    # print('\n', function_inputs.T , '\n')
    # print(new_dias.T)

    # print(solution.sum(), solution.mean())
    print(np.sum(np.abs(new_dias - function_inputs)), np.mean(np.abs(new_dias - function_inputs)))

print('\n', len(_seed_range), ' runs completed. Time taken :', time.time() - s_time)



dias_means = np.array([np.mean(np.abs(new_dias - function_inputs)) for new_dias in dia_list])
print(dias_means)

_idx = np.where(dias_means == np.min(dias_means))[0][0]
print(_idx)
new_dias = dia_list[_idx]

diameter = torch.tensor(new_dias, dtype=torch.float32)

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

print("Total change in diameters: ", dia_change.sum())
print("Mean change in diameters: ", dia_change.mean())

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_gad_" + args.wdn + "_" + str(s) + "_" + str(len(eval_samples)) + "_" + str(num_generations) + "_" + str(random_mutation_min_val) + "_" + str(random_mutation_max_val) + "_dia_dict.csv")


with open(saving_dir+"/planning_dia_gad_" + args.wdn + "_" + str(s) + "_" + str(len(eval_samples)) + "_" + str(num_generations) + "_" + str(random_mutation_min_val) + "_" + str(random_mutation_max_val) + "_wds_eval.pickle", "wb") as _file:
    pickle.dump(wds_eval, _file)
with open(saving_dir+"/planning_dia_gad_" + args.wdn + "_" + str(s) + "_" + str(len(eval_samples)) + "_" + str(num_generations) + "_" + str(random_mutation_min_val) + "_" + str(random_mutation_max_val) + "_diameter.pickle", "wb") as _file:
    pickle.dump(diameter, _file)
with open(saving_dir+"/planning_dia_gad_" + args.wdn + "_" + str(s) + "_" + str(len(eval_samples)) + "_" + str(num_generations) + "_" + str(random_mutation_min_val) + "_" + str(random_mutation_max_val) + "_dia_list.pickle", "wb") as _file:
    pickle.dump(dia_list, _file)