Mass WIP commit to reflect changes in Graph representation

src/Network.jl

@@ -1,18 +1,33 @@
+using OrderedCollections
+
 using Cairo, Compose
 using Graphs, MetaGraphs, GraphPlot
 using ModelParameters
 
-import YAML: write_file
+import YAML: load_file, write_file
 
 
 struct StreamfallNetwork
     mg::MetaDiGraph
 end
 
+function load_network(name::String, fn::String)
+    network = load_file(
+        fn;
+        dicttype=OrderedDict
+    )
+
+    return create_network(name, network)
+end
+
 
 Base.getindex(sn::StreamfallNetwork, n::String) = get_node(sn, n)
 Base.getindex(sn::StreamfallNetwork, nid::Int) = get_node(sn, nid)
 
+function node_names(sn::StreamfallNetwork)
+    verts = vertices(sn.mg)
+    return [sn[v].name for v in verts]
+end
 
 function set_prop!(sn::StreamfallNetwork, nid::Int64, prop::Symbol, var::Any)::Nothing
     MetaGraphs.set_prop!(sn.mg, nid, prop, var)
@@ -111,15 +126,15 @@ end
 
 
 """
-    create_node(mg::MetaDiGraph, node_name::String, details::Dict, nid::Int)
+    create_node(mg::MetaDiGraph, node_name::String, details::OrderedDict, nid::Int)
 
 Create a node specified with given name (if it does not exist).
 
-Returns 
-- `this_id`, ID of node (if pre-existing) and 
+Returns
+- `this_id`, ID of node (if pre-existing) and
 - `nid`, incremented node id for entire network (equal to `this_id` if exists)
 """
-function create_node(mg::MetaDiGraph, node_name::String, details::Dict, nid::Int)
+function create_node(mg::MetaDiGraph, node_name::String, details::OrderedDict, nid::Int)
     details = copy(details)
 
     match = collect(MetaGraphs.filter_vertices(mg, :name, node_name))
@@ -149,46 +164,47 @@ function create_node(mg::MetaDiGraph, node_name::String, details::Dict, nid::Int
         set_props!(mg, nid, Dict(:name=>node_name,
                                  :node=>n,
                                  :nfunc=>func))
-        
+
         this_id = nid
-        nid = nid + 1
     else
         this_id = match[1]
     end
 
-    return this_id, nid
+    return this_id
 end
 
 
 """
-    create_network(name::String, network::Dict)::StreamfallNetwork
+    create_network(name::String, network::OrderedDict)::StreamfallNetwork
 
 Create a StreamNetwork from a YAML-derived specification.
 
 # Example
 ```julia-repl
-julia> network_spec = YAML.load_file("example_network.yml")
+julia> using OrderedCollections
+julia> network_spec = YAML.load_file("example_network.yml"; dicttype=OrderedDict{Any,Any})
 julia> sn = create_network("Example Network", network_spec)
 ```
 """
-function create_network(name::String, network::Dict)::StreamfallNetwork
+function create_network(name::String, network::OrderedDict)::StreamfallNetwork
     num_nodes = length(network)
     mg = MetaDiGraph(num_nodes)
     MetaGraphs.set_prop!(mg, :description, name)
-
-    nid = 1
-    for (node, details) in network
+
+    for (nid, (node, details)) in enumerate(network)
         n_name = string(node)
 
-        this_id, nid = create_node(mg, n_name, details, nid)
+        this_id = create_node(mg, n_name, details, nid)
 
         if haskey(details, "inlets")
             inlets = details["inlets"]
-
             if !isnothing(inlets)
                 for inlet in inlets
-                    in_id, nid = create_node(mg, string(inlet), network[inlet], nid)
-                    add_edge!(mg, in_id => this_id)
+
+                    inlet_id = findall(keys(network) .== inlet)[1]
+
+                    in_id = create_node(mg, string(inlet), network[inlet], inlet_id)
+                    add_edge!(mg, inlet_id => nid)
                 end
             end
         end
@@ -200,8 +216,9 @@ function create_network(name::String, network::Dict)::StreamfallNetwork
                 @assert length(outlets) <= 1 || throw(ArgumentError(msg))
 
                 for outlet in outlets
-                    out_id, nid = create_node(mg, string(outlet), network[outlet], nid)
-                    add_edge!(mg, this_id => out_id)
+                    outlet_id = findall(keys(network) .== outlet)[1]
+                    out_id = create_node(mg, string(outlet), network[outlet], outlet_id)
+                    add_edge!(mg, nid => outlet_id)
                 end
             end
         end
@@ -303,10 +320,18 @@ function Base.show(io::IO, sn::StreamfallNetwork)
 
     vs = vertices(sn.mg)
 
-    for nid in vs
-        println(io, "Node $(nid) : \n")
+    show_verts = vs
+    if length(vs) > 4
+        show_verts = [1, 2, nothing, length(vs)-1, length(vs)]
+    end
+
+    for nid in show_verts
+        if isnothing(nid)
+            println(io, "⋮ \n")
+            continue
+        end
+        println(io, "Node $(nid): \n")
         show(io, sn[nid])
-        print("\n")
     end
 end
 
@@ -317,15 +342,15 @@ end
 Simple plot of stream network.
 """
 function plot_network(sn::StreamfallNetwork; as_html=false)
-    node_names = ["$(n.name)" for n in sn]
+    node_labels = ["$(sn[i].name)\n"*string(nameof(typeof(sn[i]))) for i in vertices(sn.mg)]
 
     if as_html
         plot_func = gplothtml
     else
         plot_func = gplot
     end
 
-    plot_func(sn.mg, nodelabel=node_names)
+    plot_func(sn.mg, nodelabel=node_labels)
 end
 
 

src/Nodes/DamNode.jl

@@ -56,33 +56,31 @@ Base.@kwdef mutable struct DamNode{P, A<:AbstractFloat} <: NetworkNode
     level::Array{A} = []
     discharge::Array{A} = []
     outflow::Array{A} = []
-
 end
 
 
 function DamNode(
     name::String,
-    area::Float64,
-    max_storage::Float64,
-    storage_coef::Float64,
-    initial_storage::Float64,
+    area::F,
+    max_storage::F,
+    storage_coef::F,
+    initial_storage::F,
     calc_dam_level::Function,
     calc_dam_area::Function,
     calc_dam_discharge::Function,
-    calc_dam_outflow::Function)
+    calc_dam_outflow::Function) where {F<:Float64}
     return DamNode(name, area, max_storage, storage_coef,
                    calc_dam_level, calc_dam_area, calc_dam_discharge, calc_dam_outflow,
-                   [initial_storage], [], [], [], [], [], [], [])
+                   F[initial_storage], F[], F[], F[], F[], F[], F[], F[])
 end
 
-
 """
     DamNode(name::String, spec::Dict)
 
 Create DamNode from a given specification.
 """
-function DamNode(name::String, spec::Dict)
-    n = DamNode{Param}(; name=name, area=spec["area"],
+function DamNode(name::String, spec::Union{Dict,OrderedDict})
+    n = DamNode{Param,Float64}(; name=name, area=spec["area"],
                        max_storage=spec["max_storage"])
 
     node_params = spec["parameters"]
@@ -127,8 +125,22 @@ function storage(node::DamNode)
     return last(node.storage)
 end
 
+function prep_state!(node::DamNode, timesteps::Int64)
+    resize!(node.storage, timesteps+1)
+    node.storage[2:end] .= 0.0
+
+    node.effective_rainfall = zeros(timesteps)
+    node.et = zeros(timesteps)
+    node.inflow = zeros(timesteps)
+    node.dam_area = zeros(timesteps)
+
+    node.level = zeros(timesteps)
+    node.discharge = zeros(timesteps)
+    node.outflow = zeros(timesteps)
+end
+
 
-function update_state(node::DamNode, storage, rainfall, et, area, discharge, outflow)
+function update_state!(node::DamNode, storage, rainfall, et, area, discharge, outflow)
     push!(node.storage, storage)
     push!(node.effective_rainfall, rainfall)
     push!(node.et, et)
@@ -139,6 +151,18 @@ function update_state(node::DamNode, storage, rainfall, et, area, discharge, out
 
     return nothing
 end
+function update_state!(node::DamNode, ts::Int64, storage, rainfall, et, area, discharge, outflow)
+    node.storage[ts+1] = storage
+
+    node.effective_rainfall[ts] = rainfall
+    node.et[ts] = et
+    node.level[ts] = node.calc_dam_level(storage)
+    node.dam_area[ts] = area
+    node.discharge[ts] = discharge
+    node.outflow[ts] = outflow
+
+    return nothing
+end
 
 
 """
@@ -172,10 +196,14 @@ end
 function run_node!(node::DamNode, climate::Climate;
                    inflow=nothing, extraction=nothing, exchange=nothing)
     timesteps = sim_length(climate)
+    prep_state!(node, timesteps)
+
     for ts in 1:timesteps
         run_node!(node, climate, ts;
                   inflow=inflow, extraction=extraction, exchange=exchange)
     end
+
+    return nothing
 end
 
 
@@ -194,23 +222,25 @@ Run a specific node for a specified time step.
 - `exchange::DataFrame` : Time series of groundwater flux
 """
 function run_node!(node::DamNode, climate::Climate, timestep::Int;
-                   inflow=nothing, extraction=nothing, exchange=nothing)
+                   inflow=nothing, extraction=nothing, exchange=nothing)::Nothing
     ts = timestep
-    if checkbounds(Bool, node.outflow, ts)
-        if node.outflow[ts] != undef
-            # already ran for this time step so no need to run
-            return node.outflow[ts], node.level[ts]
-        end
-    end
+    # if checkbounds(Bool, node.outflow, ts)
+    #     if node.outflow[ts] != undef
+    #         # already ran for this time step so no need to run
+    #         return node.outflow[ts], node.level[ts]
+    #     end
+    # end
 
     node_name = node.name
     rain, et = climate_values(node, climate, ts)
     wo = timestep_value(ts, node_name, "releases", extraction)
     ex = timestep_value(ts, node_name, "exchange", exchange)
     in_flow = timestep_value(ts, node_name, "inflow", inflow)
-    vol = node.storage[ts]
+    current_vol = node.storage[ts]
+
+    run_node!(node, ts, rain, et, current_vol, in_flow, wo, ex)
 
-    return run_node!(node, rain, et, vol, in_flow, wo, ex)
+    return nothing
 end
 
 
@@ -229,23 +259,23 @@ Calculate outflow for the dam node for a single time step.
 - outflow from dam
 """
 function run_node!(node::DamNode,
+                   ts::Int64,
                    rain::Float64,
                    et::Float64,
-                   vol::Float64,
+                   volume::Float64,
                    inflow::Float64,
                    extractions::Float64,
                    gw_flux::Float64)
-    volume = vol
     dam_area = node.calc_dam_area(volume)
     discharge = node.calc_dam_discharge(volume, node.max_storage)
 
     updated_store = update_volume(volume, inflow, gw_flux, rain, et,
                                   dam_area, extractions, discharge, node.max_storage)
     outflow = node.calc_dam_outflow(discharge, extractions)
 
-    update_state(node, updated_store, rain, et, dam_area, discharge, outflow)
+    update_state!(node, ts, updated_store, rain, et, dam_area, discharge, outflow)
 
-    return outflow, level(node)
+    return nothing
 end