Additional constructors for TwoLevelTree

docs/src/manual/multi.md

@@ -73,30 +73,46 @@ If there is uncertainty at a strategic level, this can be incorporated using the
 time structure. This structure is represented by a tree, with each node corresponding to a strategic period that contains an operational time structure.
 The operational time structure can be any combination of the *[described structures](@ref man-oper)*.
 
-The following example demonstrates how to create a regular tree  (through the function [`regular_tree`](@ref))
-where each strategic period spans 3 years and is represented by a week with daily resolution.
-The second  argument to the [`regular_tree`](@ref) function specifies the number
-of branches at each stage of the tree, excluding the first stage.
+The following example demonstrates how to create a regular tree through a constructor (`TwoLevelTree(duration::S, branching::Vector, ts::OP; op_per_strat::Float64 = 1.0) where {S,T,OP<:TimeStructure{T}}`) where each strategic period spans 3 years and is represented by a week with daily resolution.
+The second  argument of the constructor function specifies the number of branches at each stage of the tree, excluding the first stage.
+
 ```@repl ts
 using TimeStruct
-operational = SimpleTimes(7, 1);
-two_level_tree = regular_tree(3, [3,2], operational; op_per_strat = 52);
+week = SimpleTimes(7, 1);
+two_level_tree = TwoLevelTree(3, [3, 2], week; op_per_strat = 52.0);
 ```
 
 ![Illustration of TwoLevelTree](./../figures/two_level_tree.png)
 
 The branching probabilities are equal for all branches as indicated in green in the figure.
 
-!!! note "Constructors for TwoLevelTree"
-    Currently, the functionality for creating  `TwoLevelTree`'s is limited. Future versions of the package
-    will expand this functionality to allow creating trees with varying probabilities and different operational
-    time structures for the nodes.
+We also provide the possibility of having differing tree structures through the application of the [`TreeNode`](@ref) type.
+A [`TreeNode`](@ref) approach for above's time structure would be given by
 
+```@repl ts
+using TimeStruct
+week = SimpleTimes(7, 1);
+two_level_tree = TwoLevelTree(
+    TreeNode(3, week, [
+        TreeNode(3, week, 2,
+            TreeNode(3, week)
+        )
+        TreeNode(3, week, [0.5, 0.5],
+            TreeNode(3, week)
+        )
+        TreeNode(3, week, [
+            TreeNode(3, week),
+            TreeNode(3, week),
+        ])
+    ]);
+    op_per_strat = 52.0
+)
+```
+
+Similar as for [`TwoLevel`](@ref), the strategic nodes can be iterated using [`strat_periods`](@ref).
+It is possible to connect the nodes to their predecessor by iterating using the [`withprev`](@ref) iterator that returns a tuple with the parent or nothing if no parent, together with the node itself.
+This provides the flexibility to track decisions in the tree as shown by the following example that allows investment into new capacity in each strategic node while tracking the accumulated capacity.
 
-Similar as for [`TwoLevel`](@ref), the strategic nodes can be iterated using [`strat_periods`](@ref). It is possible to connect the nodes to their predecessor by
-iterating using the [`withprev`](@ref) iterator that returns a tuple with the parent or nothing if no parent, together with the node itself. This provides
-the flexibility to track decisions in the tree as shown by the following example that allows investment into new capacity in each strategic node
-while tracking the accumulated capacity.
 ```@repl ts
 using JuMP
 
@@ -110,10 +126,11 @@ for (prev, sp) in withprev(strat_pers)
     end
 end
 ```
-To ensure consistency across the tree, it is possible to iterate through all strategic scenarios
-in the tree using [`strategic_scenarios`](@ref). Here each scenario is a path from the root node
-to one of the leaves of the tree. In the example above, if we only allow one investment in the
-planning period, this can be added by restricting the number of investments in each scenario:
+
+To ensure consistency across the tree, it is possible to iterate through all strategic scenarios in the tree using [`strategic_scenarios`](@ref).
+Here each scenario is a path from the root node to one of the leaves of the tree.
+In the example above, if we only allow one investment in the planning period, this can be added by restricting the number of investments in each scenario:
+
 ```@repl ts
 for sc in strategic_scenarios(two_level_tree)
     @constraint(m, sum(invest[sp] for sp in sc) <= 1)

docs/src/reference/api.md

@@ -17,6 +17,7 @@ OperationalScenarios
 RepresentativePeriods
 TwoLevel
 TwoLevelTree
+TreeNode
 regular_tree
 ```
 

src/TimeStruct.jl

@@ -34,6 +34,8 @@ export RepresentativePeriods
 export TwoLevel
 export TwoLevelTree
 
+export TreeNode
+
 export TimeProfile
 export FixedProfile
 export OperationalProfile

src/op_scenarios/core_types.jl

@@ -12,16 +12,19 @@ and an associated probability. These scenarios are in general represented as
 [`SimpleTimes`](@ref).
 
 !!! note
-    - All scenarios must use the same type for the duration, _.i.e._, either Integer or Float.
+    - The `TimeStructure`s of all operational scenarios must use the same type for the
+      duration, *.i.e.*, either Integer or Float.
     - If the `probability` is not specified, it assigns the same probability to each scenario.
     - It is possible that `sum(probability)` is larger or smaller than 1. This can lead to
       problems in your application. Hence, it is advised to scale it. Currently, a warning
       will be given if the period shares do not sum to one as an automatic scaling will
       correspond to a breaking change.
 
 ## Example
+
 The following examples create a time structure with 2 operational scenarios corresponding to
 a single day with equal probability.
+
 ```julia
 day = SimpleTimes(24, 1)
 OperationalScenarios(2, day)

src/representative/core_types.jl

@@ -19,7 +19,8 @@ has an associated share that specifies how much of the total duration that
 is attributed to it.
 
 !!! note
-    - All representative periods must use the same type for the `TimeStructure`.
+    - The `TimeStructure`s of all representative periods must use the same type for the
+      duration, *.i.e.*, either Integer or Float.
     - If the field `period_share` is not specified, it assigns the same probability to each
       representative period.
     - It is possible that `sum(period_share)` is larger or smaller than 1. This can lead to
@@ -29,7 +30,8 @@ is attributed to it.
     - If you include [`OperationalScenarios`](@ref) in your time structure, it is important
       that the scenarios are within the representative periods, and not the other way.
 
-### Example
+## Example
+
 ```julia
 # A year represented by two days with hourly resolution and relative shares of 0.7 and 0.3
 RepresentativePeriods(8760, [0.7, 0.3], [SimpleTimes(24, 1), SimpleTimes(24,1)])

src/strat_scenarios/core_types.jl

@@ -1,23 +1,78 @@
 """
     mutable struct TwoLevelTree{S,T,OP<:AbstractTreeNode{S,T}} <: TimeStructure{T}
 
-Time structure allowing for a tree structure for the strategic level.
-
-For each strategic node in the tree a separate time structure is used for
-operational decisions. Iterating the structure will go through all operational periods.
+    TwoLevelTree(node::TreeNode; op_per_strat=8760.0)
+    TwoLevelTree(duration::S, branching::Vector, ts::OP; op_per_strat::Float64 = 1.0) where {S,T,OP<:TimeStructure{T}}
+
+Time structure allowing for a tree structure for the strategic level. For each strategic
+node in the tree a separate time structure is used for operational decisions. Iterating the
+structure will go through all operational periods.
+
+The default approach for creating a `TwoLevelTree` is by providing the root `[TreeNode`](@ref)
+with all its children nodes. In the case of a regular structure, that is all children nodes
+have the same `duration`, time structure `ts`, probability, and children itself, you can use
+a simplified constructor with the `branching` vector. The vector `branching` specifies the
+number of branchings at each stage of the tree, excluding the first stage. The branches at
+each stage will all have equal probability, duration, and time structure.
+
+!!! warning "Additional iteratores"
+    `TwoLevelTree` utilize a separate [`withprev`](@ref) method which is equivalent to the
+    existing method for the other time structures. [`withnext`](@ref), [`chunk](@ref) and
+    [`chunk_duration`](@ref) are not implemented and will result in an error when used.
+
+## Example
+
+```julia
+# Declare the individual time structure
+day = SimpleTimes(24, 1)
+
+# Regular tree with 3 strategic periods of duration 5, 3 branches for the second strategic
+# period, and 6 branchs in the thirdand forth strategic period
+regtree_1 = TwoLevelTree(5, [3, 2, 1], day)
+
+# Equivalent structure using `TreeNode` and the different constructors
+regtree_2 = TwoLevelTree(
+    TreeNode(5, day, [
+        TreeNode(5, day, 2,
+            TreeNode(5, day, TreeNode(5, day))
+        )
+        TreeNode(5, day, [0.5, 0.5],
+            TreeNode(5, day, TreeNode(5, day))
+        )
+        TreeNode(5, day, [
+            TreeNode(5, day, TreeNode(5, day)),
+            TreeNode(5, day, TreeNode(5, day)),
+        ])
+    ])
+)
+```
 """
 mutable struct TwoLevelTree{S,T,OP<:AbstractTreeNode{S,T}} <: TimeStructure{T}
     len::Int
     root::Any
     nodes::Vector{OP}
     op_per_strat::Float64
 end
+function TwoLevelTree(parent::TreeNode; op_per_strat = 1.0)
+    nodes = StratNode[]
+
+    add_node!(nodes, parent, nothing, 1.0, 1, op_per_strat)
+    nodes = convert(Array{typejoin(typeof.(nodes)...)}, nodes)
+    len = maximum([_strat_per(sn) for sn in nodes])
 
-function TwoLevelTree{S,T,OP}(
-    nodes::Vector{OP},
-    op_per_strat,
-) where {S,T,OP<:AbstractTreeNode{S,T}}
-    return TwoLevelTree{S,T,OP}(0, nothing, nodes, op_per_strat)
+    return TwoLevelTree(len, nodes[1], nodes, op_per_strat)
+end
+function TwoLevelTree(
+    duration::S,
+    branching::Vector,
+    ts::OP;
+    op_per_strat::Float64 = 1.0,
+) where {S,T,OP<:TimeStructure{T}}
+    node = TreeNode(duration, ts)
+    for k in reverse(branching)
+        node = TreeNode(duration, ts, k, node)
+    end
+    return TwoLevelTree(node; op_per_strat)
 end
 
 function _multiple_adj(itr::TwoLevelTree, n)
@@ -105,7 +160,7 @@ end
 """
     struct StrategicScenario
 
-Desription of an individual strategic scenario. It includes all strategic nodes
+Description of an individual strategic scenario. It includes all strategic nodes
 corresponding to a scenario, including the probability. It can be utilized within a
 decomposition algorithm.
 """
@@ -135,15 +190,17 @@ struct StrategicScenarios
 end
 
 """
+    strategic_scenarios(ts::TwoLevel)
     strategic_scenarios(ts::TwoLevelTree)
 
 This function returns a type for iterating through the individual strategic scenarios of a
 `TwoLevelTree`. The type of the iterator is dependent on the type of the
 input `TimeStructure`.
 
-When the `TimeStructure` is a `TimeStructure`, `strategic_scenarios` returns a
+When the `TimeStructure` is a [`TwoLevel`](@ref), `strategic_scenarios` returns a Vector with
+the `TwoLevel` as a single entry.
 """
-strategic_scenarios(two_level::TwoLevel) = [two_level]
+strategic_scenarios(ts::TwoLevel) = [ts]
 
 """
 When the `TimeStructure` is a [`TwoLevelTree`](@ref), `strategic_scenarios` returns the
@@ -162,58 +219,55 @@ function Base.iterate(scs::StrategicScenarios, state = 1)
     node = getleaf(scs.ts, state)
     prob = probability_branch(node)
     nodes = [node]
-    while !isnothing(node.parent)
-        node = node.parent
+    while !isnothing(_parent(node))
+        node = _parent(node)
         pushfirst!(nodes, node)
     end
 
     return StrategicScenario(prob, nodes), state + 1
 end
 
 """
-    add_node(
-        tree::TwoLevelTree{T, StratNode{S, T, OP}},
-        parent,
-        sp,
-        duration::S,
-        branch_prob,
-        branching,
-        oper::OP,
-    ) where {S, T, OP<:TimeStructure{T}}
-
-Iterative addition of nodes.
+    add_node!(
+        nodes::Vector{<:StratNode},
+        node::TreeNode{S, T, OP, U},
+        parent::Union{Nothing,StratNode},
+        prob::Float64,
+        sp::Int64,
+        op_per_strat::Real,
+    ) where {S,T,OP<:TimeStructure{T},U}
+
+Iterative addition of a `TreeNode` `node` to a `Vector{<:StratNode}` .
 """
-# Add nodes iteratively in a depth first manner
-function add_node(
-    tree::TwoLevelTree{S,T,StratNode{S,T,OP}},
-    parent,
-    sp,
-    duration::S,
-    branch_prob,
-    branching,
-    oper::OP,
-) where {S,T,OP<:TimeStructure{T}}
-    prob_branch = branch_prob * (isnothing(parent) ? 1.0 : parent.prob_branch)
-    mult_sp = duration * tree.op_per_strat / _total_duration(oper)
-    node = StratNode{S,T,OP}(
+# Ignored docstring, just fyi in this case
+function add_node!(
+    nodes::Vector{<:StratNode},
+    node::TreeNode{S,T,OP,U},
+    parent::Union{Nothing,StratNode},
+    prob::Float64,
+    sp::Int64,
+    op_per_strat::Real,
+) where {S,T,OP<:TimeStructure{T},U}
+    oper = node.ts
+    new_node = StratNode(
         sp,
-        branches(tree, sp) + 1,
-        duration,
-        mult_sp,
-        prob_branch,
+        count(n -> _strat_per(n) == sp, nodes) + 1,
+        duration_strat(node),
+        duration_strat(node) * op_per_strat / _total_duration(oper),
+        prob,
         parent,
         oper,
     )
-    push!(tree.nodes, node)
-    if isnothing(parent)
-        tree.root = node
-    end
+    push!(nodes, new_node)
 
-    if sp < tree.len
-        for i in 1:branching[sp]
-            # TODO: consider branching probability as input, but use uniform for now
-            add_node(tree, node, sp + 1, duration, 1.0 / branching[sp], branching, oper)
-        end
+    # Iterate through the children and add their nodes
+    for (sub_prob, sub_tn) in zip(node.probability, children(node))
+        # Continue when reaching leaf node
+        isnothing(sub_tn) && continue
+
+        # Add the new node
+        total_prob = prob * sub_prob
+        add_node!(nodes, sub_tn, new_node, total_prob, sp + 1, op_per_strat)
     end
 end
 
@@ -223,25 +277,32 @@ end
         branching::Vector,
         ts::OP;
         op_per_strat::Real=1.0,
-    ) where {S, T, OP<:TimeStructure{T}}
+    ) where {S,T,OP<:TimeStructure{T}}
 
 Function for creating a regular tree with a uniform structure for each strategic period.
+
 Each strategic period is of equal length as given by `duration` and will have the same
 operational time structure `ts`. The vector `branching` specifies the number of branchings
 at each stage of the tree, excluding the first stage. The branches at each stage will
 all have equal probability.
+
+!!! note "Deprecated function"
+    This function is deprecated and will be removed in a later release. The new function is
+    given by
+
+    ```julia
+    TwoLevelTree(duration, branching, ts; op_per_strat)
+    ```
+
 """
 function regular_tree(
     duration::S,
     branching::Vector,
     ts::OP;
     op_per_strat::Real = 1.0,
 ) where {S,T,OP<:TimeStructure{T}}
-    tree = TwoLevelTree{S,T,StratNode{S,T,OP}}(Vector{StratNode{S,T,OP}}(), op_per_strat)
-    tree.len = length(branching) + 1
-    add_node(tree, nothing, 1, duration, 1.0, branching, ts)
-
-    return tree
+    op_per_strat = convert(Float64, op_per_strat)
+    return TwoLevelTree(duration, branching, ts; op_per_strat)
 end
 
 """