dataset_generator.py

# -*- coding: utf-8 -*-
""" 
    Using partial code from 
    Vrachimis et al. https://github.com/KIOS-Research/BattLeDIM
"""
import pandas as pd
import numpy as np
import wntr
import pickle
import os
import argparse
import time
from math import sqrt, ceil
import warnings, copy, json
warnings.filterwarnings('ignore')
import matplotlib.pyplot as plt

# -*- coding: utf-8 -*-
"""
Created on Sun Apr 22 15:58:48 2018

@author: mkiria01
"""
import scipy.io
from utils.utils import create_graph, WDN_Graph
import torch



def create_cli_parser():
    # ----- ----- ----- ----- ----- -----
    # Command line arguments
    # ----- ----- ----- ----- ----- -----
    parser  = argparse.ArgumentParser()
    parser.add_argument('--wdn',
                        default = 'l_town',
                        type    = str,
                        choices = ['hanoi', 'fossolo', 'pescara', 'l_town', 'zhijiang'],
                        help    = "specify the WDS for which you want to simulate the scenarios; default is l_town ")
    parser.add_argument('--start_scenario',
                        default = 1,
                        type    = int,
                        help    = "specify the start scenario name, must be an integer; default is 1")
    parser.add_argument('--end_scenario',
                        default = 50,
                        type    = int,
                        help    = "specify the end scenario name, must be an integer; default is 50")
    parser.add_argument('--mu_dem',
                        default = 1.0,
                        type    = float,
                        help    = "Specify the mean of the normal ditribution used to sample toy demand patterns; default is 1.0.")   
    parser.add_argument('--sigma_dem',
                        default = 0.15,
                        type    = float,
                        help    = "Specify the mean of the normal ditribution used to sample toy demand patterns; default is 0.15.")   
    parser.add_argument('--sigma_dia',
                        default = 1/30,
                        type    = float,
                        help    = "Specify the standard deviation of the noise to be added to the diameters; default is 1/30")   
    parser.add_argument('--_seed',
                        default = None,
                        type    = int,
                        help    = "Specify the random seed for noise; default is None, where it will be set to the scenario name for every scenario.")   
    parser.add_argument('--start_time',
                        default = '2018-01-01 00:00',
                        type    = str,
                        help    = "Specify the start time of the simulation; default is 2018-01-01 00:00, the samples will be sampled every 30 minutes starting from this time.")   
    parser.add_argument('--end_time',
                        default = '2018-01-14 23:30',
                        type    = str,
                        help    = "Specify the end time of the simulation; default is 2018-01-14 23:30.")   
    parser.add_argument('--scenario_type',
                        default = 'toy',
                        type    = str,
                        help    = "choose scenario type from 'toy' and 'real'; default is 'toy'.")
    parser.add_argument('--pattern_gen',
                        default = False,
                        action  = 'store_true',
                        help    = "specify if realistic demand patterns need to be generated; default is False.")
    parser.add_argument('--dem_multiplier',
                        default = 0.1,
                        type    = float,
                        help    = "specify demand multiplier in case of realistic demand patterns (only effective when --pattern_gen=True).")
    parser.add_argument('--pattern_multiplier',
                        default = 0.5,
                        type    = float,
                        help    = "specify pattern multiplier in case of realistic demand patterns (only effective when --pattern_gen=True).")
    parser.add_argument('--qunc_multiplier',
                        default = 0.5,
                        type    = float,
                        help    = "specify scaling factor for the offsets sampling range for demands and pipe attributes \
                                    in case of realistic demand patterns (only effective when --pattern_gen=True).")

    return parser



def genDem():
    """
        A method to generate realistic demands for different WDS.
    """
    weekPat = scipy.io.loadmat('weekPat_30min.mat')
    Aw = weekPat['Aw']
    nw = weekPat['nw']
    yearOffset = scipy.io.loadmat('yearOffset_30min.mat')
    Ay = yearOffset['Ay']
    ny = yearOffset['ny']
     
    # Create yearly component
    days = 365
    
    T=(288/6)*days # one year period in five minute intervals
    w=2*np.pi/T
    k=np.arange(1, days*288/6+1 ,1) # number of time steps in time series
    n=ny[0][0] # number of fourier coefficients
    Hy=[1]*len(k)
    
    for i in range(1,n+1):
        Hy=np.column_stack((Hy, np.sin(i*w*k), np.cos(i*w*k)))
    
    Hy.shape # check size matrix
    uncY=0.1
    AyR = Ay*(1-uncY+2*uncY*np.random.rand(int(Ay.shape[0]),int(Ay.shape[1]))) # randomize fourier coefficients
    yearOffset = np.dot(Hy, AyR)
    
    # Create weekly component
    T=(288/6)*7 #one week period in five minute intervals
    w=2*np.pi/T
    k=np.arange(1, days*288/6+1 ,1) # number of time steps in time series
    n=nw[0][0] # number of fourier coefficients
    Hw=[1]*len(k)
    for i in range(1,n+1):
        Hw=np.column_stack((Hw, np.sin(i*w*k), np.cos(i*w*k)))
    
    uncW=0.1
    AwR = Aw*(1-uncW+2*uncW*np.random.rand(int(Aw.shape[0]),int(Aw.shape[1]))) # randomize fourier coefficients
    weekYearPat = np.dot(Hw, AwR)
    
    # Create random component
    uncR=0.05
    random = np.random.normal(0,(-uncR+2*uncR),(int(weekYearPat.shape[0]),int(weekYearPat.shape[1]))) #normally distributed random numbers
    
    # Create demand
    #blow=30
    #bhigh=35
    base =1#blow+np.random.rand()*(bhigh-blow)
    variation = 0.75+ np.random.normal(0,0.07) # from 0 to 1
    dem = base * (yearOffset+1) * (weekYearPat*variation+1) * (random+1)
    dem = dem.tolist()
    demFinal = []
    for d in dem:
        demFinal.append(d[0])
      
    return demFinal



class DatasetCreator:
    def __init__(self, wdn, scenario_folder, inp_file, start_time, end_time,
                 qunc=np.arange(0, 0.25, 0.05), dem_multiplier=1., pattern_multiplier=1.,
                 pattern_gen=False, _print=True, return_heads=False, wdn_object=None):

        self.scenario_folder = scenario_folder
        if _print:
            print(f'Run input file: "{inp_file}"')
        self.results_folder = os.path.join(os.getcwd(), scenario_folder, "Results-Clean")
        # Create Results folder
        if not os.path.isdir(self.results_folder):
            os.makedirs(self.results_folder)

        # demand-driven (DD) or pressure dependent demand (PDD)
        Mode_Simulation = 'DD'  

        # Load EPANET network file
        if wdn_object is None:
            self.wn = wntr.network.WaterNetworkModel(inp_file)
        else:
            self.wn = wdn_object
        self.wn.options.hydraulic.demand_model = Mode_Simulation

        self.nodes = self.wn.get_graph().nodes()
        self.links = self.wn.link_name_list

        # Get time step
        self.time_step = round(self.wn.options.time.hydraulic_timestep)
        # Create time_stamp
        self.time_stamp = pd.date_range(start_time, end_time, freq=str(self.time_step / 60) + "min")

        # Simulation duration in steps
        self.wn.options.time.duration = (len(self.time_stamp) - 1) * self.time_step  

        if pattern_gen:
            # For Demand Pattern Generation
            qunc_index = int(round(np.random.uniform(len(qunc)-1)))
            uncertainty_Length = qunc[qunc_index]
            
            qunc_index = int(round(np.random.uniform(len(qunc)-1)))
            uncertainty_Diameter = qunc[qunc_index]
            
            qunc_index = int(round(np.random.uniform(len(qunc)-1)))
            uncertainty_Roughness = qunc[qunc_index]
            
            qunc_index = int(round(np.random.uniform(len(qunc)-1)))
            uncertainty_base_demand = qunc[qunc_index]

            ###########################################################################  
            ## SET BASE DEMANDS AND PATTERNS      
            # Remove all patterns
            #        # Initial base demands SET ALL EQUAL 1
            #if "ky" in INP:
            self.wn._patterns= {}
            tempbase_demand = self.wn.query_node_attribute('base_demand')
            tempbase_demand = np.array([tempbase_demand[i] for i, line in enumerate(tempbase_demand)]) * dem_multiplier
            
            tmp = list(map(lambda x: x * uncertainty_base_demand, tempbase_demand))
            ql=tempbase_demand-tmp
            qu=tempbase_demand+tmp
            mtempbase_demand=len(tempbase_demand)
            qext_mtempbase_demand=ql+np.random.rand(mtempbase_demand)*(qu-ql)
            
            for w, junction in enumerate(self.wn.junction_name_list):
                self.wn.get_node(junction).demand_timeseries_list[0].base_value = qext_mtempbase_demand[w] #self.wn.query_node_attribute('base_demand')
                pattern_name = 'P_'+junction
                patts = genDem()
                patts = list(((np.array(patts) - 1) * pattern_multiplier) + 1)
                self.wn.add_pattern(pattern_name, patts)
                
                for patterns in self.wn.nodes._data[junction].demand_timeseries_list._list:
                    patterns.pattern_name = pattern_name
                    patterns.pattern.name = pattern_name
            
            ###########################################################################
            ## SET UNCERTAINTY PARAMETER
            # Uncertainty Length
            tempLengths = self.wn.query_link_attribute('length')
            tempLengths = np.array([tempLengths[i] for i, line in enumerate(tempLengths)])
            tmp = list(map(lambda x: x * uncertainty_Length, tempLengths))
            ql=tempLengths-tmp
            qu=tempLengths+tmp
            mlength=len(tempLengths)
            qext=ql+np.random.rand(mlength)*(qu-ql)
                
            # Uncertainty Diameter
            tempDiameters = self.wn.query_link_attribute('diameter')
            tempDiameters = np.array([tempDiameters[i] for i, line in enumerate(tempDiameters)])
            tmp = list(map(lambda x: x * uncertainty_Diameter, tempDiameters))
            ql=tempDiameters-tmp
            qu=tempDiameters+tmp
            dem_diameter=len(tempDiameters)
            diameters=ql+np.random.rand(dem_diameter)*(qu-ql)
                
            # Uncertainty Roughness
            tempRoughness = self.wn.query_link_attribute('roughness')
            tempRoughness = np.array([tempRoughness[i] for i, line in enumerate(tempRoughness)])
            tmp = list(map(lambda x: x * uncertainty_Roughness, tempRoughness))
            ql=tempRoughness-tmp
            qu=tempRoughness+tmp
            dem_roughness=len(tempRoughness)
            qextR=ql+np.random.rand(dem_roughness)*(qu-ql)
            for w, line1 in enumerate(qextR):
                self.wn.get_link(self.wn.link_name_list[w]).roughness=line1
                self.wn.get_link(self.wn.link_name_list[w]).length=qext[w]
                self.wn.get_link(self.wn.link_name_list[w]).diameter=diameters[w]

            filename = os.path.join(scenario_folder, wdn + ".inp")
            out_inp = wntr.epanet.io.InpFile()
            out_inp.write(
                filename, 
                wn = self.wn,
                units = None, 
                version = 2.2, 
                force_coordinates = False
                )
            inp_file = filename      
    

    def dataset_generator(self, scenario_times=[], _print=True, return_heads=False):
        # Path of EPANET Input File
        if _print:
            print(f"Dataset Generator run...")
        
        if not return_heads:
            # Save the water network model to a file before using it in a simulation
            with open(os.path.join(os.getcwd(), self.scenario_folder,'self.wn.pickle'), 'wb') as f:
                pickle.dump(self.wn, f)

        # Run wntr simulator
        scenario_start_time = time.time()
        sim = wntr.sim.WNTRSimulator(self.wn)
        results = sim.run_sim()
        scenario_end_time = time.time()
        scenario_time = scenario_end_time - scenario_start_time
        scenario_times.append(scenario_time)
        if _print:
            print('... simulation done!')
        
        if results.node["pressure"].empty:
            print("Negative pressures.")
            return -1

        if results:
            decimal_size = 6            

            # Create xlsx file with Measurements
            def export_measurements(pressure_sensors, flow_sensors, file_out="Measurements.xlsx", return_heads=False):
                total_pressures = {'Timestamp': self.time_stamp}
                total_demands = {'Timestamp': self.time_stamp}
                total_flows = {'Timestamp': self.time_stamp}
                total_levels = {'Timestamp': self.time_stamp}
                total_heads = {'Timestamp': self.time_stamp}
                for j in range(0, self.wn.num_nodes):
                    node_id = self.wn.node_name_list[j]
                    if node_id in pressure_sensors:
                        pres = results.node['pressure'][node_id]
                        pres = pres[:len(self.time_stamp)]
                        pres = [round(elem, decimal_size) for elem in pres]
                        total_pressures[node_id] = pres

                        head = results.node['head'][node_id]
                        head = head[:len(self.time_stamp)]
                        head = [round(elem, decimal_size) for elem in head]
                        total_heads[node_id] = head

                        dem = results.node['demand'][node_id]
                        dem = dem[:len(self.time_stamp)]
                        dem = [round(elem, decimal_size) for elem in dem] 
                        total_demands[node_id] = dem

                        level_pres = results.node['pressure'][node_id]
                        level_pres = level_pres[:len(self.time_stamp)]
                        level_pres = [round(elem, decimal_size) for elem in level_pres]
                        total_levels[node_id] = level_pres

                for j in range(0, self.wn.num_links):
                    link_id = self.wn.link_name_list[j]

                    if link_id not in flow_sensors:
                        continue
                    flows = results.link['flowrate'][link_id]
                    flows = [round(elem, decimal_size) for elem in flows]
                    flows = flows[:len(self.time_stamp)]
                    total_flows[link_id] = flows

                # Loading original demands
                dem_multiplier = self.wn.options.hydraulic.demand_multiplier
                n_timesteps = len(self.time_stamp)
                orig_demands = {'Timestamp': self.time_stamp}
                for node_id in self.wn.node_name_list:
                    node_dem = 0
                    if self.wn.nodes._data[node_id].node_type == 'Junction':
                        for patterns in self.wn.nodes._data[node_id].demand_timeseries_list._list:
                            if patterns.pattern.multipliers is not None:
                                node_dem += (patterns.base_value * dem_multiplier) * patterns.pattern.multipliers[: n_timesteps]                                    
                            else:
                                node_dem += patterns.base_value * dem_multiplier
                    try:
                        repeat_idx = ceil(n_timesteps / len(node_dem))
                        node_dem_copy = copy.deepcopy(node_dem)
                        for i in range(1, repeat_idx):
                            node_dem = np.concatenate((node_dem, node_dem_copy))
                        orig_demands[node_id] = node_dem[: n_timesteps] #* 3600 * 1000
                    except:
                        orig_demands[node_id] = node_dem              


                # Create a Pandas dataframe from the data.                
                df1 = pd.DataFrame(total_pressures)
                df2 = pd.DataFrame(total_demands)
                df3 = pd.DataFrame(total_flows)
                df4 = pd.DataFrame(total_levels)
                df5 = pd.DataFrame(total_heads)
                df6 = pd.DataFrame(orig_demands)

                if return_heads:
                    return df5.values[:, 1:]

                if _print:
                    print("Minimum Pressure: ", df1[df1 != 0].min(numeric_only=True).min(), "Maximum Pressure: ", df1.max(numeric_only=True).max())
                    print("Minimum Head: ", df5.min(numeric_only=True).min(), "Maximum Head: ", df5.max(numeric_only=True).max())

                # Create a Pandas Excel writer using XlsxWriter as the engine.
                writer = pd.ExcelWriter(os.path.join(self.results_folder, file_out), engine='xlsxwriter')

                # Convert the dataframe to an XlsxWriter Excel object.
                # Pressures (m), Demands (m^3/s), Flows (m^3/s), Levels (m), Heads (m)
                df1.to_excel(writer, sheet_name='Pressures (m)', index=False)
                df2.to_excel(writer, sheet_name='Demands (m3_s)', index=False)
                df3.to_excel(writer, sheet_name='Flows (m3_s)', index=False)
                df4.to_excel(writer, sheet_name='Levels (m)', index=False)
                df5.to_excel(writer, sheet_name='Heads (m)', index=False)
                df6.to_excel(writer, sheet_name='Orig_Demands (m3_s)', index=False)

                # Close the Pandas Excel writer and output the Excel file.
                writer._save()

                # Export as .csv files -- .csv files are much faster parsed by pandas than huge .xlsx files!
                df1.to_csv(os.path.join(self.results_folder, file_out.replace(".xlsx", "_Pressures.csv")), index=False)
                df2.to_csv(os.path.join(self.results_folder, file_out.replace(".xlsx", "_Demands.csv")), index=False)
                df3.to_csv(os.path.join(self.results_folder, file_out.replace(".xlsx", "_Flows.csv")), index=False)
                df4.to_csv(os.path.join(self.results_folder, file_out.replace(".xlsx", "_Levels.csv")), index=False)
                df5.to_csv(os.path.join(self.results_folder, file_out.replace(".xlsx", "_Heads.csv")), index=False)
                df6.to_csv(os.path.join(self.results_folder, file_out.replace(".xlsx", "_Orig_Demands.csv")), index=False)

            # Export all measurements 
            if return_heads:
                heads = export_measurements(self.nodes, self.links, "Measurements_All.xlsx", return_heads=return_heads)
            else:
                export_measurements(self.nodes, self.links, "Measurements_All.xlsx")
                # Clean up
                os.remove(os.path.join(os.getcwd(), self.scenario_folder,'self.wn.pickle'))
        else:
            print('Results empty.')
            return -1
        
        if return_heads:
            return heads
        
        return scenario_times
    


def run(wdn = 'l_town', start_scenario=1, end_scenario=50, scenario_times=[], 
        pattern_gen=False, qunc=np.arange(0, 0.25, 0.05), pattern_multiplier=1.0,
        mu_dem=1.0, sigma_dem=0.1, sigma_dia=1/30, scenario_type="toy", dem_multiplier=1.0,
        in_seed=None, start_time='2018-01-01 00:00', end_time='2018-01-14 23:30'):

    for s in range(start_scenario, end_scenario + 1):
        scenario = 's' + str(s)
        
        save_dir = os.path.join(os.getcwd(),"networks",  wdn, scenario_type)

        scenario_dir = os.path.join(save_dir, scenario)
        if not os.path.isdir(scenario_dir):
            os.makedirs(scenario_dir)

        if wdn == "hanoi" and s < 1000:
            inp_file = os.path.join(scenario_dir, "Hanoi_CMH_Scenario-" + str(s) + ".inp")
        else:
            inp_file = os.path.join(save_dir, wdn + ".inp")

        wn = wntr.network.WaterNetworkModel(inp_file)
        if in_seed is None:
            _seed = s
        else:
            _seed = in_seed
        np.random.seed(_seed)
        print('Seed used for scenario ', s)

        t = time.time()
    
        if not pattern_gen and s >= 1000:
            _len = 48
            for node_id in wn.node_name_list:
                if wn.nodes._data[node_id].node_type == 'Junction' and 'leak' not in node_id:
                    for patterns in wn.nodes._data[node_id].demand_timeseries_list._list:
                        pattern = np.round(np.random.normal(mu_dem, sigma_dem, size = _len), 6).clip(0) 
                        pattern_offset = np.round(np.random.normal(0, sigma_dem, size = _len), 6).clip(0)  
                        wn.add_pattern(
                            name = "random_week_"+str(node_id), 
                            pattern = pattern + pattern_offset
                            )
                        patterns.pattern_name = "random_week_"+str(node_id)
                        patterns.pattern.name = "random_week_"+str(node_id)
                        patterns.pattern.multipliers = pattern + pattern_offset

            for key, value in wn.links._data.items():
                if wdn == "fossolo":
                    wn.links._data[key].diameter = wn.links._data[key].diameter * 2 

                wn.links._data[key].diameter = \
                    wn.links._data[key].diameter * ( 1 + np.random.normal(0, sigma_dia, size=1)[0] ) 
            
            filename = os.path.join(scenario_dir, wdn + ".inp")
            out_inp = wntr.epanet.io.InpFile()
            out_inp.write(
                filename, 
                wn = wn,
                units = None, 
                version = 2.2, 
                force_coordinates = False
                )
            inp_file = filename

        if wdn == "fossolo":
            for key, value in wn.links._data.items():
                wn.links._data[key].diameter = wn.links._data[key].diameter * 2 
            filename = os.path.join(scenario_dir, wdn + ".inp")
            out_inp = wntr.epanet.io.InpFile()
            out_inp.write(
                filename, 
                wn = wn,
                units = None, 
                version = 2.2, 
                force_coordinates = False
                )
            inp_file = filename

        # Call dataset creator        
        L = DatasetCreator(wdn, scenario_dir, inp_file, start_time, end_time, 
                           pattern_gen=pattern_gen, qunc=qunc, pattern_multiplier=pattern_multiplier, dem_multiplier=dem_multiplier)
        scenario_times = L.dataset_generator(scenario_times)

        print('\nScenario ' + scenario + ' generated. Total Elapsed time is ' + str(time.time() - t) + ' seconds.\n')

    return scenario_times




def simulate_wntr(args, s=2000, s_offset=2001, diameter=None, dem_multiplier=None, demand_series=None, dem_offset=None,
                   _simulate=True, start_time="2018-01-01 00:00", end_time="2018-12-31 23:30", 
                   _print=True, save_data=True):
    '''
    Run EPANET on the wds specified in args.inp_file. Loads default configuration from scenario data specified by s and s_offset.
    Set wds attributes as provided in <diameter>, <demand_series>.
    Alternatively offset demand values by <dem_offset>, or scale them via <dem_multiplier>.
    '''

    wn = wntr.network.WaterNetworkModel(args.inp_file)


    if demand_series is not None:
        for idx, node_id in enumerate(wn.node_name_list):
            if wn.nodes._data[node_id].node_type == 'Junction':
                wn.nodes._data[node_id].demand_timeseries_list[0].base_value = demand_series[0, idx].item()
                for patterns in wn.nodes._data[node_id].demand_timeseries_list._list:            
                    patterns.pattern.multipliers = list(np.ones_like(patterns.pattern.multipliers))
                    pattern = (demand_series[:, idx] / demand_series[0, idx]).flatten()
                    pattern = torch.nan_to_num(pattern, nan=0, posinf=0, neginf=0).numpy()
                    wn.add_pattern(
                        name = "eval_samples_"+str(node_id), 
                        pattern = pattern 
                        )
                    patterns.pattern_name = "eval_samples_"+str(node_id)
                    patterns.pattern.multipliers = pattern 

    if dem_offset is not None:
        for w, junction in enumerate(wn.junction_name_list):
            wn.get_node(junction).demand_timeseries_list[0].base_value = \
                wn.get_node(junction).demand_timeseries_list[0].base_value + dem_offset[w]

    if dem_multiplier is not None:
        for w, junction in enumerate(wn.junction_name_list):
            wn.get_node(junction).demand_timeseries_list[0].base_value = \
                wn.get_node(junction).demand_timeseries_list[0].base_value * dem_multiplier

    if diameter is not None:    
        for idx, (key, value) in zip(np.arange(len(wn.links._data)), wn.links._data.items()):
            wn.links._data[key].diameter = diameter.clone().numpy()[idx,0] 

    s_n = s + s_offset

    if args.toy:
        scenario_path = os.path.join(os.getcwd(),"networks",  args.wdn, "toy", "s"+str(s_n))
    elif args.real:
        scenario_path = os.path.join(os.getcwd(),"networks",  args.wdn, "real", "s"+str(s_n))

    if save_data:

        inp_file = os.path.join(scenario_path, args.wdn + ".inp")

        if not os.path.isdir(scenario_path):
            os.makedirs(scenario_path)

        out_inp = wntr.epanet.io.InpFile()
        out_inp.write(
            inp_file, 
            wn = wn,
            units = None, 
            version = 2.2, 
            force_coordinates = False
            )

        if _simulate:

            # Call dataset creator        
            L = DatasetCreator(args.wdn, 
                            scenario_path, 
                            inp_file,
                            start_time=start_time, 
                            end_time=end_time, 
                            pattern_gen=False,
                            _print=_print
                            )
            scenario_times = L.dataset_generator(scenario_times=[], _print=_print)

        args.path_to_data = os.path.join(scenario_path, "Results-Clean", "Measurements_All.xlsx")

        wds_recon, reservoirs = create_graph(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_demands = torch.tensor(wntr_d_df.astype("float32").values)

        if _print:
            print('\t', str(s_n))
            print(wds_recon.X.shape, wds_recon.edge_indices[0].shape, wds_recon.edge_attr[0].shape)

        return wds_recon, reservoirs, wntr_demands    
    
    if _simulate and not save_data:
        # Call dataset creator        
        L = DatasetCreator(
                        args.wdn, 
                        scenario_path, 
                        args.inp_file,
                        start_time=start_time, 
                        end_time=end_time, 
                        pattern_gen=False,
                        _print=_print,
                        return_heads=True,
                        wdn_object=wn,
                        )
        heads = L.dataset_generator(scenario_times=[], _print=_print, return_heads=True)

        return heads
    
    print('Please either specify _simulate or save_data')
    return None




if __name__ == '__main__':
    parser = create_cli_parser()
    args = parser.parse_args()  
    args = vars(args)
    with open('args_dg.json', 'r') as args_file:
        args.update(json.load(args_file))
    args['toy'] = args.pop('scenario_type') == 'toy'
    args['real'] = args.pop('scenario_type') == 'real'
    assert args['real'] != args['toy'] # XOR

    args = argparse.Namespace(**args)

    # hanoi         mu_dem, sigma_dem, sigma_dia,       1.0, 0.15, 1/30
    # fossolo       mu_dem, sigma_dem, sigma_dia,       1.0, 0.15, 1/30     dem_multiplier 3        dia 2 times
    # pescara       mu_dem, sigma_dem, sigma_dia,       1.0, 0.15, 1/30     dem_multiplier 0.5
    # l_town        mu_dem, sigma_dem, sigma_dia,       1.0, 0.15, 1/30     dem_multiplier 5
    # zhijiang      mu_dem, sigma_dem, sigma_dia,       1.0, 0.15, 1/30     dem_multiplier 0.1

    
    qunc = np.arange(0, .25, .05) 


    # fossolo       dem_multiplier      3.      pattern_multiplier        .5        qunc_multiplier     .5     dia 2 times
    # pescara       dem_multiplier      .5      pattern_multiplier        .5        qunc_multiplier     .5
    # Area-C        dem_multiplier      5.      pattern_multiplier        .5        qunc_multiplier     .5
    # zhijang       dem_multiplier      .1      pattern_multiplier        .5        qunc_multiplier     .5

    
    """ For planning diameters """
    # dem_multiplier, pattern_multiplier, qunc_multiplier 
    # 2000          hanoi (1.5, .5, .5)         fossolo (6., .5, .5)           pescara (.85, .5, .5)        l_town (8.5, .5, .5)          zhijiang (.11, .5, .5)

    qunc = qunc * args.qunc_multiplier


    scenario_times = []
    scenario_times = run(
                        wdn=args.wdn, 
                        start_scenario=args.start_scenario, 
                        end_scenario=args.end_scenario, 
                        scenario_times=scenario_times,
                        mu_dem=args.mu_dem, 
                        sigma_dem=args.sigma_dem, 
                        sigma_dia=args.sigma_dia, 
                        in_seed=args._seed, 
                        start_time=args.start_time, 
                        end_time=args.end_time,
                        pattern_gen=args.pattern_gen, 
                        qunc=qunc, 
                        dem_multiplier=args.dem_multiplier,
                        pattern_multiplier=args.pattern_multiplier,
                        scenario_type = args.scenario_type
                        )    

    print("Total Simulation Time for all Scenarios: ", np.sum(scenario_times))
    print("Scenarios Simulation Times: ", scenario_times)