Improvement of utilities for TwoLevelTree

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)

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

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))

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))")

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,45 +1104,75 @@ 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]
     @test TimeStruct._opscen(op) == 1
     @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
 
     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],