diff --git a/docs/src/reference/api.md b/docs/src/reference/api.md index ed89c9e..418494d 100644 --- a/docs/src/reference/api.md +++ b/docs/src/reference/api.md @@ -28,6 +28,16 @@ mult_scen mult_repr mult_strat probability_scen +probability_branch +``` + +The following properties are included for a [`TwoLevelTree`](@ref) time structure: + +```@docs +n_strat_per +n_children +n_leaves +n_branches ``` ## [Iterating time structures](@id api-iter) diff --git a/src/TimeStruct.jl b/src/TimeStruct.jl index b190f20..52a8d58 100644 --- a/src/TimeStruct.jl +++ b/src/TimeStruct.jl @@ -55,6 +55,7 @@ export multiple_strat export mult_scen, mult_repr, mult_strat export start_time, end_time, remaining export start_oper_time, end_oper_time +export n_strat_per, n_children, n_leaves, n_branches export Discounter export discount diff --git a/src/strat_scenarios/core_types.jl b/src/strat_scenarios/core_types.jl index 7003a96..7a96750 100644 --- a/src/strat_scenarios/core_types.jl +++ b/src/strat_scenarios/core_types.jl @@ -75,22 +75,58 @@ function TwoLevelTree( return TwoLevelTree(node; op_per_strat) end -function _multiple_adj(itr::TwoLevelTree, n) +function _multiple_adj(itr::TwoLevelTree, n::Int) mult = itr.nodes[n].duration * itr.op_per_strat / _total_duration(itr.nodes[n].operational) return stripunit(mult) end strat_nodes(ts::TwoLevelTree) = ts.nodes -function children(n::StratNode, ts::TwoLevelTree) - return [c for c in ts.nodes if _parent(c) == n] -end -nchildren(n::StratNode, ts::TwoLevelTree) = count(c -> _parent(c) == n, strat_nodes(ts)) +""" + n_strat_per(ts::TwoLevelTree) + +Returns the number of strategic periods of a [`TwoLevelTree`](@ref). If the number is +different for the individual branches, it returns the maximum value. +""" +n_strat_per(ts::TwoLevelTree) = maximum(_strat_per(c) for c in strat_nodes(ts)) + +""" + n_children(n::StratNode, ts::TwoLevelTree) + +Returns the number of children of a [`StratNode`](@ref). +""" +n_children(n::StratNode, ts::TwoLevelTree) = count(c -> _parent(c) == n, strat_nodes(ts)) + +""" + n_leaves(ts::TwoLevelTree) + +Returns the number of children of a [`TwoLevelTree`](@ref). +""" +n_leaves(ts::TwoLevelTree) = count(n -> n_children(n, ts) == 0, strat_nodes(ts)) + +""" + n_branches(ts::TwoLevelTree, sp::Int) -branches(ts::TwoLevelTree, sp::Int) = count(n -> _strat_per(n) == sp, strat_nodes(ts)) -leaves(ts::TwoLevelTree) = [n for n in strat_nodes(ts) if nchildren(n, ts) == 0] -nleaves(ts::TwoLevelTree) = count(n -> nchildren(n, ts) == 0, strat_nodes(ts)) -getleaf(ts::TwoLevelTree, leaf::Int) = leaves(ts)[leaf] +Returns the number of branches in strategic period `sp` of a [`TwoLevelTree`](@ref). +""" +n_branches(ts::TwoLevelTree, sp::Int) = count(n -> _strat_per(n) == sp, strat_nodes(ts)) + +children(n::StratNode, ts::TwoLevelTree) = [c for c in ts.nodes if _parent(c) == n] +leaves(ts::TwoLevelTree) = [n for n in strat_nodes(ts) if n_children(n, ts) == 0] + +get_leaf(ts::TwoLevelTree, leaf::Int) = leaves(ts)[leaf] +function get_strat_node(ts::TwoLevelTree, sp::Int, branch::Int) + node = filter(n -> _strat_per(n) == sp && _branch(n) == branch, strat_nodes(ts)) + if isempty(node) + throw( + ErrorException( + "The `TwoLevelTree` does not have a node with strategic period $(sp) and branch $(branch)", + ), + ) + else + return node[1] + end +end function Base.length(itr::TwoLevelTree) return sum(length(n.operational) for n in itr.nodes) @@ -210,13 +246,13 @@ strategic_scenarios(ts::TwoLevelTree) = StrategicScenarios(ts) # Allow a TwoLevel structure to be used as a tree with one scenario # TODO: Should be replaced with a single wrapper as it is the case for the other scenarios -Base.length(scens::StrategicScenarios) = nleaves(scens.ts) +Base.length(scens::StrategicScenarios) = n_leaves(scens.ts) function Base.iterate(scs::StrategicScenarios, state = 1) - if state > nleaves(scs.ts) + if state > n_leaves(scs.ts) return nothing end - node = getleaf(scs.ts, state) + node = get_leaf(scs.ts, state) prob = probability_branch(node) nodes = [node] while !isnothing(_parent(node)) diff --git a/src/structures.jl b/src/structures.jl index 3e2abf1..dc7b0da 100644 --- a/src/structures.jl +++ b/src/structures.jl @@ -84,10 +84,18 @@ multiple(t::TimePeriod) = 1.0 """ probability(t::TimePeriod) + Returns the probability associated with the time period. """ probability(t::TimePeriod) = 1.0 +""" + probability_branch(t::Union{TimePeriod, TimeStructurePeriod}) + +Returns the branch probability associated with the time period or time structure period +""" +probability_branch(t::Union{TimePeriod,TimeStructurePeriod}) = 1.0 + # Functions used for indexing into time profiles # TODO: Consider either setting all as default to one, including _oper, or none _oper(t::TimePeriod) = error("_oper() not implemented for $(typeof(t))") diff --git a/test/runtests.jl b/test/runtests.jl index 397c68c..1d045f2 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -369,11 +369,14 @@ end @test pers[1] < pers[2] @test pers[24] < pers[25] - sp = collect(strat_periods(ts)) + sps = collect(strat_periods(ts)) # Test that collect is working correctly - ops = collect(sp[1]) - @test sum(ops[it] == op for (it, op) in enumerate(sp[1])) == 24 + ops = collect(sps[1]) + @test sum(ops[it] == op for (it, op) in enumerate(sps[1])) == 24 + + # Test that the branch probability is 1 + @test all(probability_branch(sp) == 1.0 for sp in sps) end @testitem "TwoLevel with units" begin @@ -1093,7 +1096,6 @@ end for oscs in oscs_inv[2:end] @test [probability_branch(osc) for osc in oscs] == p_branch[1] end - return ops_inv[1] end end @@ -1102,6 +1104,9 @@ end regtree = TwoLevelTree(5, [3, 2], SimpleTimes(5, 1)) n_sp = 10 n_op = n_sp * 5 + + @test all(TimeStruct._multiple_adj(regtree, t) == 1 for t in 1:n_sp) + ops = TwoLevelTreeTest.fun(regtree, n_sp, n_op) op = ops[31] @@ -1109,30 +1114,57 @@ end @test TimeStruct._strat_per(op) == 3 @test TimeStruct._branch(op) == 4 @test TimeStruct._oper(op) == 1 + @test TimeStruct._rper(op) == 1 @test duration(op) == 1 @test probability(op) == 1 / 6 @test op isa eltype(typeof(regtree)) + @test repr(op) == "sp3-br4-t1" + @test op < ops[32] + @test ops[4] < op + @test TimeStruct.StrategicTreeIndexable(typeof(op)) == TimeStruct.HasStratTreeIndex() nodes = strat_nodes(regtree) for sp in 1:3 - @test sum(TimeStruct.probability_branch(n) for n in nodes if n.sp == sp) ≈ 1.0 + @test sum(probability_branch(n) for n in nodes if n.sp == sp) ≈ 1.0 end node = nodes[2] @test length(node) == 5 @test first(node) isa eltype(typeof(node)) + @test repr(node) == "sp2-br1" + @test all(multiple_strat(node, t) ≈ 0.2 for t in node) + @test TimeStruct.StrategicTreeIndexable(typeof(node)) == TimeStruct.HasStratTreeIndex() + # Test helper functions leaves = TimeStruct.leaves(regtree) - @test length(leaves) == TimeStruct.nleaves(regtree) - - scens = collect(TimeStruct.strategic_scenarios(regtree)) + @test length(leaves) == n_leaves(regtree) + @test leaves[3] == TimeStruct.get_leaf(regtree, 3) + @test n_branches(regtree, 2) == 3 + @test n_branches(regtree, 3) == 6 + + @test n_children(regtree.root, regtree) == 3 + @test TimeStruct.children(regtree.root, regtree) == regtree.nodes[[2, 5, 8]] + @test TimeStruct.get_strat_node(regtree, 2, 1) == regtree.nodes[2] + @test TimeStruct.get_strat_node(regtree, 3, 2) == regtree.nodes[4] + @test TimeStruct.get_strat_node(regtree, 3, 3) == regtree.nodes[6] + @test TimeStruct.get_strat_node(regtree, 3, 6) == regtree.nodes[10] + @test_throws ErrorException TimeStruct.get_strat_node(regtree, 3, 8) + @test_throws ErrorException TimeStruct.get_strat_node(regtree, 4, 1) + + @test n_strat_per(regtree) == 3 + + # Test strategic periods + sps = strat_periods(regtree) + @test eltype(sps) == typeof(node) + @test sps[2] == node + + # Test strategic scenarios + scens = collect(strategic_scenarios(regtree)) + @test length(strategic_scenarios(regtree)) == 6 @test length(scens[2].nodes) == regtree.len + @test last(scens[1]) == regtree.nodes[3] @test scens[3].nodes[1] == regtree.nodes[1] - ssp = TimeStruct.StrategicStochasticProfile([ - [10], - [11, 12, 13], - [20, 21, 22, 23, 30, 40], - ]) + ssp = StrategicStochasticProfile([[10], [11, 12, 13], [20, 21, 22, 23, 30, 40]]) @test ssp[nodes[3]] == 20 @test ssp[nodes[8]] == 13 @@ -1140,7 +1172,7 @@ end price1 = OperationalProfile([1, 2, 2, 5, 6]) price2 = FixedProfile(4) - dsp = TimeStruct.StrategicStochasticProfile([ + dsp = StrategicStochasticProfile([ [price1], [price1, price2, price2], [price1, price2, price2, price1, price2, price2],