From 8d3d79c46a1f2cfced2d0cbf0dc6d37d82d293de Mon Sep 17 00:00:00 2001
From: Julian Straus <julian.straus@sintef.no>
Date: Tue, 27 Aug 2024 10:59:16 +0200
Subject: [PATCH 01/14] Initial commit: - Rework on the TwoLevelTree type -
 Works with out RepresentativePeriods or OperationalScenarios

---
 src/TimeStruct.jl   |   6 +-
 src/tree_nodes.jl   |  41 +++++
 src/twoleveltree.jl | 380 ++++++++++++++++++++------------------------
 test/runtests.jl    |  20 +--
 4 files changed, 231 insertions(+), 216 deletions(-)
 create mode 100644 src/tree_nodes.jl

diff --git a/src/TimeStruct.jl b/src/TimeStruct.jl
index a9d68a9..c6ece44 100644
--- a/src/TimeStruct.jl
+++ b/src/TimeStruct.jl
@@ -12,7 +12,6 @@ include("calendar.jl")
 include("representative.jl")
 include("stochastic.jl")
 include("twolevel.jl")
-include("twoleveltree.jl")
 include("strat_periods.jl")
 include("repr_periods.jl")
 include("opscenarios.jl")
@@ -20,6 +19,9 @@ include("profiles.jl")
 include("utils.jl")
 include("discount.jl")
 
+include("tree_nodes.jl")
+include("twoleveltree.jl")
+
 export TimeStructure
 export SimpleTimes
 export CalendarTimes
@@ -40,7 +42,7 @@ export RepresentativeProfile
 export opscenarios
 export repr_periods
 export strat_periods, strategic_periods
-export regular_tree, strat_nodes, scenarios
+export regular_tree, strat_nodes, strategic_scenarios
 export withprev, chunk, chunk_duration
 export isfirst, duration, duration_strat, multiple, probability
 export multiple_strat
diff --git a/src/tree_nodes.jl b/src/tree_nodes.jl
new file mode 100644
index 0000000..6fc9f82
--- /dev/null
+++ b/src/tree_nodes.jl
@@ -0,0 +1,41 @@
+"""
+    AbstractTreeNode{T} <: TimeStructure{T}
+
+Abstract base type for all tree nodes within a [`TwoLevelTree`] type
+"""
+abstract type AbstractTreeNode{T} <: TimeStructure{T} end
+
+struct StratNode{S, T, OP<:TimeStructure{T}} <: AbstractTreeNode{T}
+    sp::Int
+    branch::Int
+    duration::S
+    prob_branch::Float64
+    mult_sp::Float64
+    parent::Any
+    operational::OP
+end
+
+Base.show(io::IO, n::StratNode) = print(io, "sp$(n.sp)-br$(n.branch)")
+_branch(n::StratNode) = n.branch
+_strat_per(n::StratNode) = n.sp
+probability_branch(n::StratNode) = n.prob_branch
+duration_strat(n::StratNode) = n.duration
+
+isfirst(n::StratNode) = n.sp == 1
+
+# Iterate through time periods of a strategic node
+Base.length(n::StratNode) = length(n.operational)
+Base.eltype(::Type{StratNode{T}}) where {T} = TreePeriod
+function Base.iterate(itr::StratNode, state = nothing)
+    next =
+        isnothing(state) ? iterate(itr.operational) :
+        iterate(itr.operational, state)
+    next === nothing && return nothing
+    per = next[1]
+
+    mult = itr.mult_sp * multiple(per)
+    return TreePeriod(itr.sp, itr.branch, probability_branch(itr), mult, per),
+    next[2]
+end
+
+multiple_strat(sp::StratNode, t) = multiple(t) / duration_strat(sp)
diff --git a/src/twoleveltree.jl b/src/twoleveltree.jl
index 83758a8..c7fe3b4 100644
--- a/src/twoleveltree.jl
+++ b/src/twoleveltree.jl
@@ -1,275 +1,196 @@
 
-abstract type AbstractTreeNode{T} end
-
 """
     mutable struct TwoLevelTree{T} <: TimeStructure{T}
 
-    Time structure allowing for a tree structure for 
-    the strategic level. 
+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.
+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.
 """
-mutable struct TwoLevelTree{T} <: TimeStructure{T}
+mutable struct TwoLevelTree{T, OP<:AbstractTreeNode{T}} <: TimeStructure{T}
     len::Int
     root::Any
-    nodes::Vector{<:AbstractTreeNode{T}}
+    nodes::Vector{OP}
+    op_per_strat::Float64
+end
+
+function TwoLevelTree{T, OP}(nodes::Vector{OP}, op_per_strat) where {T, OP<:AbstractTreeNode{T}}
+    return TwoLevelTree{T, OP}(0, nothing, nodes, op_per_strat)
+end
+
+function Base.length(itr::TwoLevelTree)
+    return sum(length(n.operational) for n in itr.nodes)
 end
+Base.eltype(::Type{TwoLevelTree{T,OP}}) where {T,OP} = TreePeriod
 
-function TwoLevelTree{T}(nodes::Vector{<:AbstractTreeNode}) where {T}
-    return TwoLevelTree{T}(0, nothing, nodes)
+function _multiple_adj(itr::TwoLevelTree, n)
+    mult =
+        itr.nodes[n].duration * itr.op_per_strat /
+        _total_duration(itr.nodes[n].operational)
+    return stripunit(mult)
 end
+strat_nodes(tree::TwoLevelTree) = tree.nodes
+
+children(n::StratNode, ts::TwoLevelTree) = [c for c in ts.nodes if c.parent == n]
+nchildren(n::StratNode, ts::TwoLevelTree) = count(c -> c.parent == n, ts.nodes)
+
+branches(tree::TwoLevelTree, sp) = count(n -> n.sp == sp, tree.nodes)
+
+leaves(ts::TwoLevelTree) = [n for n in ts.nodes if nchildren(n, ts) == 0]
+nleaves(ts::TwoLevelTree) = count(n -> nchildren(n, ts) == 0, ts.nodes)
+getleaf(ts::TwoLevelTree, leaf) = leaves(ts)[leaf]
 
 """
-	struct OperPeriod <: TimePeriod
+	struct TreePeriod{P} <: TimePeriod where {P<:TimePeriod}
+
+Time period for iteration of a `TwoLevelTree` time structure. This period has in addition
+to an operational period also the two fields `branch` and `prob_branch` corresponding to the
+respective branch and probability of the branch
 
-    Time period for iteration of a TwoLevelTree time structure. 
+!!! warn "Using OperationalScenarios"
+    The probability will always only correspond to the branch probability, even when you
+    utilize `OperationalScenarios`.  Using the function `probability` includes however the
+    scenario probability.
 """
-struct OperPeriod <: TimePeriod
+struct TreePeriod{P} <: TimePeriod where {P<:TimePeriod}
     sp::Int
     branch::Int
-    sc::Int
-    op::Int
-    duration::Float64
-    prob::Float64
+    prob_branch::Float64
+    multiple::Float64
+    period::P
 end
 
-duration(t::OperPeriod) = t.duration
-probability(t::OperPeriod) = t.prob
-_oper(t::OperPeriod) = t.op
-_opscen(t::OperPeriod) = t.sc
-_strat_per(t::OperPeriod) = t.sp
-_branch(t::OperPeriod) = t.branch
+isfirst(t::TreePeriod) = isfirst(t.period)
+duration(t::TreePeriod) = duration(t.period)
+probability_branch(t::TreePeriod) = t.prob_branch
+probability(t::TreePeriod) = probability(t.period) * probability_branch(t)
+multiple(t::TreePeriod) = t.multiple
 
-function Base.length(itr::TwoLevelTree)
-    return sum(length(n.strat_node.operational) for n in itr.nodes)
+_oper(t::TreePeriod) = _oper(t.period)
+_opscen(t::TreePeriod) = _opscen(t.period)
+_rper(t::TreePeriod) = _rper(t.period)
+_branch(t::TreePeriod) = t.branch
+_strat_per(t::TreePeriod) = t.sp
+
+function Base.show(io::IO, t::TreePeriod)
+    return print(io, "sp$(t.sp)-br$(t.branch)-$(t.period)")
+end
+function Base.isless(t1::TreePeriod, t2::TreePeriod)
+    return t1.period < t2.period
 end
-Base.eltype(::Type{TwoLevelTree{T}}) where {T} = OperPeriod
 
 # Iterate through all time periods as OperationalPeriods
 function Base.iterate(itr::TwoLevelTree)
-    spn = itr.nodes[1].strat_node
+    spn = itr.nodes[1]
     next = iterate(spn.operational)
     next === nothing && return nothing
     per = next[1]
-    return OperPeriod(
-        spn.sp,
-        spn.branch,
-        _opscen(per),
-        _oper(per),
-        duration(per),
-        probability(spn) * probability(per),
-    ),
+
+    mult = _multiple_adj(itr, 1) * multiple(per)
+    return TreePeriod(spn.sp, spn.branch, probability_branch(spn), mult, per),
     (1, next[2])
 end
 
 function Base.iterate(itr::TwoLevelTree, state)
     i = state[1]
-    spn = itr.nodes[i].strat_node
+    spn = itr.nodes[i]
     next = iterate(spn.operational, state[2])
     if next === nothing
         i = i + 1
         if i > length(itr.nodes)
             return nothing
         end
-        spn = itr.nodes[i].strat_node
+        spn = itr.nodes[i]
         next = iterate(spn.operational)
     end
     per = next[1]
-    return OperPeriod(
-        spn.sp,
-        spn.branch,
-        _opscen(per),
-        _oper(per),
-        duration(per),
-        probability(spn) * probability(per),
-    ),
-    (i, next[2])
-end
 
-struct StratNode{T} <: TimePeriod
-    sp::Int
-    branch::Int
-    duration::Duration
-    probability::Float64
-    operational::TimeStructure{T}
-    time_struct::TimeStructure
-end
-
-Base.show(io::IO, n::StratNode) = print(io, "sp$(n.sp)-br$(n.branch)")
-_branch(n::StratNode) = n.branch
-_strat_per(n::StratNode) = n.sp
-probability(n::StratNode) = n.probability
-duration(n::StratNode) = n.duration
-
-isfirst(n::StratNode) = n.sp == 1
-
-struct TreeNode{T} <: AbstractTreeNode{T}
-    node::Int
-    parent::Union{Nothing,TreeNode}
-    strat_node::StratNode{T}
+    mult = _multiple_adj(itr, i) * multiple(per)
+    return TreePeriod(spn.sp, spn.branch, probability_branch(spn), mult, per),
+    (i, next[2])
 end
 
-children(n::TreeNode, ts::TwoLevelTree) = [c for c in ts.nodes if c.parent == n]
-nchildren(n::TreeNode, ts::TwoLevelTree) = count(c -> c.parent == n, ts.nodes)
-strat_nodes(ts::TwoLevelTree) = [n.strat_node for n in ts.nodes]
-strat_periods(ts::TwoLevelTree) = strat_nodes(ts)
-
-# Iterate through time periods of a strategic node
-Base.length(n::StratNode) = length(n.operational)
-Base.eltype(::Type{StratNode{T}}) where {T} = OperPeriod
-function Base.iterate(itr::StratNode, state = nothing)
-    next =
-        isnothing(state) ? iterate(itr.operational) :
-        iterate(itr.operational, state)
-    next === nothing && return nothing
-    per = next[1]
-    return OperPeriod(
-        itr.sp,
-        itr.branch,
-        _opscen(per),
-        _oper(per),
-        duration(per),
-        itr.probability * probability(per),
-    ),
-    next[2]
-end
+"""
+    struct StrategicScenario
 
-struct StratNodeOperationalScenario{T} <: TimeStructure{T}
-    node::StratNode
-    scen::Int
+Desription 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.
+"""
+struct StrategicScenario
     probability::Float64
-    operational::TimeStructure{T}
-end
-
-# Iteration through all time periods in an operational scenario of a strategic node
-Base.length(snops::StratNodeOperationalScenario) = length(snops.operational)
-Base.eltype(_::StratNodeOperationalScenario) = OperPeriod
-function Base.iterate(snops::StratNodeOperationalScenario, state = nothing)
-    next =
-        isnothing(state) ? iterate(snops.operational) :
-        iterate(snops.operational, state)
-    isnothing(next) && return nothing
-    per = next[1]
-    return OperPeriod(
-        snops.node.sp,
-        snops.node.branch,
-        snops.scen,
-        _oper(per),
-        duration(per),
-        snops.probability,
-    ),
-    next[2]
+    nodes::Vector{<:StratNode}
 end
 
-# Iterate through operational scenarios of a strategic node
-struct StratNodeOpScens{TS}
-    node::StratNode
-    operational::TS
-end
-
-opscenarios(n::StratNode) = StratNodeOpScens(n, n.operational)
-
-function Base.length(
-    sops::StratNodeOpScens{OperationalScenarios{T,OP}},
-) where {T,OP}
-    return sops.operational.len
-end
-function Base.iterate(
-    sops::StratNodeOpScens{OperationalScenarios{T,OP}},
-    state = 1,
-) where {T,OP}
-    state > sops.operational.len && return nothing
-    return StratNodeOperationalScenario(
-        sops.node,
-        state,
-        sops.node.probability * sops.operational.probability[state],
-        sops.operational.scenarios[state],
-    ),
-    state + 1
-end
+# Iterate through strategic periods of scenario
+Base.length(scen::StrategicScenario) = length(scen.nodes)
+Base.last(scen::StrategicScenario) = last(scen.nodes)
 
-function Base.length(sops::StratNodeOpScens{SimpleTimes{T}}) where {T}
-    return 1
-end
-function Base.iterate(
-    sops::StratNodeOpScens{SimpleTimes{T}},
-    state = 1,
-) where {T}
-    state > 1 && return nothing
-    return StratNodeOperationalScenario(
-        sops.node,
-        state,
-        sops.node.probability,
-        sops.operational,
-    ),
-    state + 1
+function Base.iterate(scs::StrategicScenario, state = nothing)
+    next = isnothing(state) ? iterate(scs.nodes) : iterate(scs.nodes, state)
+    isnothing(next) && return nothing
+    return next[1].strat_node, next[2]
 end
 
-struct Scenario
-    probability::Float64
-    nodes::Vector{TreeNode}
-end
+"""
+    struct StrategicScenarios
 
-# Iterate through all scenarios
-struct Scenarios
+Type for iteration through the individual strategic scenarios represented as
+[`StrategicScenario`](@ref).
+"""
+struct StrategicScenarios
     ts::TwoLevelTree
 end
+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
+strategic_scenarios(two_level::TwoLevel{S,T}) where {S,T} = [two_level]
 
-scenarios(ts::TwoLevelTree) = Scenarios(ts)
-
-nleaves(ts::TwoLevelTree) = count(n -> nchildren(n, ts) == 0, ts.nodes)
-leaves(ts::TwoLevelTree) = [n for n in ts.nodes if nchildren(n, ts) == 0]
-getleaf(ts::TwoLevelTree, leaf) = leaves(ts)[leaf]
 
-Base.length(scens::Scenarios) = nleaves(scens.ts)
-function Base.iterate(scs::Scenarios, state = 1)
+Base.length(scens::StrategicScenarios) = nleaves(scens.ts)
+function Base.iterate(scs::StrategicScenarios, state = 1)
     if state > nleaves(scs.ts)
         return nothing
     end
 
     node = getleaf(scs.ts, state)
-    prob = probability(node.strat_node)
+    prob = probability_branch(node)
     nodes = [node]
     while !isnothing(node.parent)
         node = node.parent
         pushfirst!(nodes, node)
     end
 
-    return Scenario(prob, nodes), state + 1
-end
-
-# Iterate through strategic periods of scenario
-Base.length(scen::Scenario) = length(scen.nodes)
-function Base.iterate(scs::Scenario, state = nothing)
-    next = isnothing(state) ? iterate(scs.nodes) : iterate(scs.nodes, state)
-    isnothing(next) && return nothing
-    return next[1].strat_node, next[2]
+    return StrategicScenario(prob, nodes), state + 1
 end
 
-branches(tree::TwoLevelTree, sp) = count(n -> n.strat_node.sp == sp, tree.nodes)
-
-# Allow a TwoLevel structure to be used as a tree with one scenario
-scenarios(two_level::TwoLevel{S,T}) where {S,T} = [two_level]
-
+"""
+    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 nodes iteratively in a depth first manner
 function add_node(
-    tree::TwoLevelTree{T},
+    tree::TwoLevelTree{T, StratNode{S, T, OP}},
     parent,
-    index,
     sp,
-    duration,
+    duration::S,
     branch_prob,
     branching,
-    ts::TimeStructure{T},
-) where {T}
-    prob =
-        branch_prob * (isnothing(parent) ? 1.0 : parent.strat_node.probability)
-    node = TreeNode{T}(
-        index,
-        parent,
-        StratNode{T}(sp, branches(tree, sp) + 1, duration, prob, ts, tree),
-    )
+    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}(sp, branches(tree, sp) + 1, duration, prob_branch, mult_sp, parent, oper)
     push!(tree.nodes, node)
     if isnothing(parent)
         tree.root = node
@@ -281,26 +202,75 @@ function add_node(
             add_node(
                 tree,
                 node,
-                length(tree.nodes) + 1,
                 sp + 1,
                 duration,
                 1.0 / branching[sp],
                 branching,
-                ts,
+                oper,
             )
         end
     end
 end
 
-# Create a regular tree with the given branching structure and the same time structure in each node 
+"""
+    regular_tree(
+        duration::S,
+        branching::Vector,
+        ts::OP;
+        op_per_strat::Real=1.0,
+    ) where {S, T, OP<:TimeStructure{T}}
+
+Function for creating a regular tree.
+"""
 function regular_tree(
-    duration,
+    duration::S,
     branching::Vector,
-    ts::TimeStructure{T},
-) where {T}
-    tree = TwoLevelTree{T}(Vector{TreeNode{T}}())
+    ts::OP;
+    op_per_strat::Real=1.0,
+) where {S, T, OP<:TimeStructure{T}}
+    tree = TwoLevelTree{T, StratNode{S, T, OP}}(Vector{StratNode{S, T, OP}}(), op_per_strat)
     tree.len = length(branching) + 1
-    add_node(tree, nothing, 1, 1, duration, 1.0, branching, ts)
+    add_node(tree, nothing, 1, duration, 1.0, branching, ts)
 
     return tree
 end
+
+"""
+    struct StratTreeNodes{T, OP}
+
+Type for iterating through the individual strategic nodes of a [`TwoLevelTree`](@ref).
+It is automatically created through the function [`strat_periods`](@ref), and hence,
+[`strategic_periods`](@ref).
+
+Iterating through `StratTreeNodes` using the WithPrev iterator changes the behaviour,
+although the meaining remains unchanged.
+"""
+struct StratTreeNodes{T, OP}
+    ts::TwoLevelTree{T, OP}
+end
+strat_periods(ts::TwoLevelTree) = StratTreeNodes(ts)
+Base.length(sps::StratTreeNodes) = length(sps.ts.nodes)
+
+function Base.iterate(stps::StratTreeNodes, state = nothing)
+    next = isnothing(state) ? 1 : state + 1
+    next == length(stps) + 1 && return nothing
+    return stps.ts.nodes[next], next
+end
+# function Base.iterate(stps::StratTreeNodes)
+#     return stps.ts.nodes[state], 1
+# end
+# function Base.iterate(stps::StratTreeNodes, state)
+#     return stps.ts.nodes[state + 1], state + 1
+# end
+
+function Base.iterate(w::WithPrev{StratTreeNodes{T,OP}}) where {T, OP}
+    n = iterate(w.itr)
+    n === nothing && return n
+    return (nothing, n[1]), (n[1], n[2])
+end
+
+function Base.iterate(w::WithPrev{StratTreeNodes{T,OP}}, state) where {T, OP}
+    n = iterate(w.itr, state[2])
+    n === nothing && return n
+    return (n[1].parent, n[1]), (n[1], n[2])
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index bc3befa..b6f70d5 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -884,6 +884,8 @@ end
     @test length(ops) == 5 * 10
     ops2 = collect(regtree)
     @test ops == ops2
+    ops3 = [t for sp in strategic_periods(regtree) for t in sp]
+    @test ops == ops3
 
     op = ops[31]
     @test TimeStruct._opscen(op) == 1
@@ -892,20 +894,20 @@ end
     @test TimeStruct._oper(op) == 1
     @test duration(op) == 1
     @test probability(op) == 1 / 6
-    @test typeof(op) == eltype(typeof(regtree))
+    @test op isa eltype(typeof(regtree))
 
     nodes = strat_nodes(regtree)
     for sp in 1:3
-        @test sum(probability(n) for n in nodes if n.sp == sp) ≈ 1.0
+        @test sum(TimeStruct.probability_branch(n) for n in nodes if n.sp == sp) ≈ 1.0
     end
     node = nodes[2]
     @test length(node) == 5
-    @test typeof(first(node)) == eltype(typeof(node))
+    @test first(node) isa eltype(typeof(node))
 
     leaves = TimeStruct.leaves(regtree)
     @test length(leaves) == TimeStruct.nleaves(regtree)
 
-    scens = collect(TimeStruct.scenarios(regtree))
+    scens = collect(TimeStruct.strategic_scenarios(regtree))
     @test length(scens[2].nodes) == regtree.len
     @test scens[3].nodes[1] == regtree.nodes[1]
 
@@ -1060,16 +1062,16 @@ end
     ) == 50
 end
 
-@testitem "Strategic scenarios" begin
+@testitem "Strategic scenarios with operational scenarios" begin
     regtree = TimeStruct.regular_tree(
         5,
         [3, 2],
         OperationalScenarios(3, SimpleTimes(5, 1)),
     )
 
-    @test length(scenarios(regtree)) == 6
+    @test length(strategic_scenarios(regtree)) == 6
 
-    for sc in scenarios(regtree)
+    for sc in strategic_scenarios(regtree)
         @test length(sc) == length(collect(sc))
 
         for (prev_sp, sp) in withprev(sc)
@@ -1084,9 +1086,9 @@ end
 @testitem "TwoLevel as a tree" begin
     two_level = TwoLevel(5, 10, SimpleTimes(10, 1))
 
-    scens = scenarios(two_level)
+    scens = TimeStruct.strategic_scenarios(two_level)
     @test length(scens) == 1
-    sps = collect(sp for sc in scenarios(two_level) for sp in strat_periods(sc))
+    sps = collect(sp for sc in TimeStruct.strategic_scenarios(two_level) for sp in strat_periods(sc))
     @test length(sps) == 5
 end
 

From f38e4553cf8235fa2f44ae6a905b95fd03c7245c Mon Sep 17 00:00:00 2001
From: Julian Straus <julian.straus@sintef.no>
Date: Tue, 27 Aug 2024 14:10:26 +0200
Subject: [PATCH 02/14] Rework of strategic StrategicStochasticProfile

---
 src/TimeStruct.jl |  4 ++--
 src/profiles.jl   | 50 +++++++++++++++++++++++++++++++++++++++--------
 src/tree_nodes.jl | 10 ++++++++++
 test/runtests.jl  |  2 +-
 4 files changed, 55 insertions(+), 11 deletions(-)

diff --git a/src/TimeStruct.jl b/src/TimeStruct.jl
index c6ece44..d1f327e 100644
--- a/src/TimeStruct.jl
+++ b/src/TimeStruct.jl
@@ -15,13 +15,14 @@ include("twolevel.jl")
 include("strat_periods.jl")
 include("repr_periods.jl")
 include("opscenarios.jl")
-include("profiles.jl")
 include("utils.jl")
 include("discount.jl")
 
 include("tree_nodes.jl")
 include("twoleveltree.jl")
 
+include("profiles.jl")
+
 export TimeStructure
 export SimpleTimes
 export CalendarTimes
@@ -36,7 +37,6 @@ export OperationalProfile
 export ScenarioProfile
 export StrategicProfile
 export StrategicStochasticProfile
-export DynamicStochasticProfile
 export RepresentativeProfile
 
 export opscenarios
diff --git a/src/profiles.jl b/src/profiles.jl
index b85a56a..124a0dd 100644
--- a/src/profiles.jl
+++ b/src/profiles.jl
@@ -165,18 +165,52 @@ function Base.getindex(
     return _value_lookup(RepresentativeIndexable(T), rp, period)
 end
 
-struct StrategicStochasticProfile{T} <: TimeProfile{T}
-    vals::Vector{Vector{T}}
+
+"""
+    StrategicStochasticProfile(vals::Vector{<:Vector{P}}) where {T<:Duration, P<:TimeProfile{T}}
+    StrategicStochasticProfile(vals::Vector{<:Vector{<:Number}})
+
+Time profile with a separate time profile for each strategic node in a [`TwoLevelTree`](@ref)
+structure.
+
+If too few profiles are provided, the last given profile will be repeated, both for strategic
+periods and branches within a strategic period.
+
+## Example
+```julia
+ # The same value in each strategic period and branch
+profile = StrategicStochasticProfile([[1], [21, 22]])
+# Varying values in each strategic period and branch
+profile = StrategicStochasticProfile([
+    [OperationalProfile([11, 12])],
+    [OperationalProfile([21, 22]), OperationalProfile([31, 32])]
+])
+```
+"""
+struct StrategicStochasticProfile{T<:Duration,P<:TimeProfile{T}} <: TimeProfile{T}
+    vals::Vector{<:Vector{P}}
 end
-function Base.getindex(ssp::StrategicStochasticProfile, i::TimePeriod)
-    return ssp.vals[_strat_per(i)][_branch(i)]
+function StrategicStochasticProfile(vals::Vector{<:Vector{<:Duration}})
+    return StrategicStochasticProfile([[FixedProfile(v_2) for v_2 in v_1] for v_1 in vals])
 end
 
-struct DynamicStochasticProfile{T<:Number} <: TimeProfile{T}
-    vals::Vector{<:Vector{<:TimeProfile{T}}}
+function _value_lookup(::HasStratTreeIndex, ssp::StrategicStochasticProfile, period)
+    sp_prof = ssp.vals[_strat_per(period) > length(ssp.vals) ? end : _strat_per(period)]
+    branch_prof = sp_prof[_branch(period) > length(sp_prof) ? end : _branch(period)]
+    return branch_prof[period]
 end
-function Base.getindex(ssp::DynamicStochasticProfile, i::TimePeriod)
-    return ssp.vals[_strat_per(i)][_branch(i)][i]
+
+function _value_lookup(::NoStratTreeIndex, ssp::StrategicStochasticProfile, period)
+    return error(
+        "Type $(typeof(period)) can not be used as index for a strategic stochastic profile",
+    )
+end
+
+function Base.getindex(
+    ssp::StrategicStochasticProfile,
+    period::T,
+) where {T<:Union{TimePeriod,TimeStructure}}
+    return _value_lookup(StrategicTreeIndexable(T), ssp, period)
 end
 
 Base.getindex(TP::TimeProfile, inds...) = [TP[i] for i in inds]
diff --git a/src/tree_nodes.jl b/src/tree_nodes.jl
index 6fc9f82..7855f11 100644
--- a/src/tree_nodes.jl
+++ b/src/tree_nodes.jl
@@ -39,3 +39,13 @@ function Base.iterate(itr::StratNode, state = nothing)
 end
 
 multiple_strat(sp::StratNode, t) = multiple(t) / duration_strat(sp)
+
+abstract type StrategicTreeIndexable end
+
+struct HasStratTreeIndex <: StrategicTreeIndexable end
+struct NoStratTreeIndex <: StrategicTreeIndexable end
+
+StrategicTreeIndexable(::Type) = NoStratTreeIndex()
+StrategicTreeIndexable(::Type{<:AbstractTreeNode}) = HasStratTreeIndex()
+StrategicTreeIndexable(::Type{<:TimePeriod}) = HasStratTreeIndex()
+StrategicIndexable(::Type{<:AbstractTreeNode}) = HasStratIndex()
diff --git a/test/runtests.jl b/test/runtests.jl
index b6f70d5..6fd9919 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -923,7 +923,7 @@ end
     price1 = OperationalProfile([1, 2, 2, 5, 6])
     price2 = FixedProfile(4)
 
-    dsp = TimeStruct.DynamicStochasticProfile([
+    dsp = TimeStruct.StrategicStochasticProfile([
         [price1],
         [price1, price2, price2],
         [price1, price2, price2, price1, price2, price2],

From 509571555dec2bd53095b7f556a92690d9fc8d49 Mon Sep 17 00:00:00 2001
From: Julian Straus <julian.straus@sintef.no>
Date: Tue, 27 Aug 2024 14:23:08 +0200
Subject: [PATCH 03/14] Added support for OperationalScenarios

---
 src/tree_nodes.jl   | 115 +++++++++++++++++++++++++++++++++++++-
 src/twoleveltree.jl |  15 ++++-
 test/runtests.jl    | 132 ++++++++++++++++++++++----------------------
 3 files changed, 193 insertions(+), 69 deletions(-)

diff --git a/src/tree_nodes.jl b/src/tree_nodes.jl
index 7855f11..60c8fc6 100644
--- a/src/tree_nodes.jl
+++ b/src/tree_nodes.jl
@@ -20,6 +20,7 @@ _branch(n::StratNode) = n.branch
 _strat_per(n::StratNode) = n.sp
 probability_branch(n::StratNode) = n.prob_branch
 duration_strat(n::StratNode) = n.duration
+multiple_strat(sp::StratNode, t) = multiple(t) / duration_strat(sp)
 
 isfirst(n::StratNode) = n.sp == 1
 
@@ -38,8 +39,6 @@ function Base.iterate(itr::StratNode, state = nothing)
     next[2]
 end
 
-multiple_strat(sp::StratNode, t) = multiple(t) / duration_strat(sp)
-
 abstract type StrategicTreeIndexable end
 
 struct HasStratTreeIndex <: StrategicTreeIndexable end
@@ -49,3 +48,115 @@ StrategicTreeIndexable(::Type) = NoStratTreeIndex()
 StrategicTreeIndexable(::Type{<:AbstractTreeNode}) = HasStratTreeIndex()
 StrategicTreeIndexable(::Type{<:TimePeriod}) = HasStratTreeIndex()
 StrategicIndexable(::Type{<:AbstractTreeNode}) = HasStratIndex()
+
+
+"""
+    struct StratNodeOperationalScenario{T} <: AbstractOperationalScenario{T}
+
+A structure representing a single operational scenario for a strategic node supporting
+iteration over its time periods.
+"""
+struct StratNodeOperationalScenario{T} <: AbstractOperationalScenario{T}
+    sp::Int
+    branch::Int
+    scen::Int
+    mult_sp::Float64
+    mult_scen::Float64
+    prob_branch::Float64
+    prob_scen::Float64
+    operational::TimeStructure{T}
+end
+
+function Base.show(io::IO, os::StratNodeOperationalScenario)
+    return print(io, "sp$(os.sp)-br$(os.branch)-sc$(os.scen)")
+end
+
+probability(os::StratNodeOperationalScenario) = os.prob_branch * prob_scen
+mult_scen(os::StratNodeOperationalScenario) = os.mult_scen
+_opscen(os::StratNodeOperationalScenario) = os.scen
+_branch(os::StratNodeOperationalScenario) = os.branch
+_strat_per(os::StratNodeOperationalScenario) = os.sp
+
+StrategicTreeIndexable(::Type{<:StratNodeOperationalScenario}) = HasStratTreeIndex()
+StrategicIndexable(::Type{<:StratNodeOperationalScenario}) = HasStratIndex()
+
+Base.length(snops::StratNodeOperationalScenario) = length(snops.operational)
+Base.eltype(_::StratNodeOperationalScenario) = TreePeriod
+
+function Base.last(os::StratNodeOperationalScenario)
+    per = last(os.operational)
+    return OperationalPeriod(os.sp, per, os.mult_sp * multiple(per))
+end
+
+function Base.getindex(os::StratNodeOperationalScenario, index)
+    per = os.operational[index]
+    mult = os.mult_sp * multiple(per)
+    return TreePeriod(os.sp, os.branch, probability_branch(os), mult, per)
+end
+
+function Base.eachindex(os::StratNodeOperationalScenario)
+    return eachindex(os.operational)
+end
+
+# Iterate the time periods of a StratOperationalScenario
+function Base.iterate(os::StratNodeOperationalScenario, state = nothing)
+    next =
+        isnothing(state) ? iterate(os.operational) :
+        iterate(os.operational, state)
+    isnothing(next) && return nothing
+
+    return TreePeriod(os.sp, os.branch, os.prob_branch, os.mult_sp * multiple(next[1]), next[1]),
+    next[2]
+end
+
+
+# Iteration through scenarios
+"""
+    struct StratNodeOpScens
+
+Type for iterating through the individual operational scenarios of a [`StratNode`](@ref).
+It is automatically created through the function [`opscenarios`](@ref).
+"""
+struct StratNodeOpScens
+    sp::Int
+    branch::Int
+    mult_sp::Float64
+    prob_branch::Float64
+    opscens::Any
+end
+
+function StratNodeOpScens(n::StratNode{S,T,OP}, opscens) where {S,T,OP<:TimeStructure{T}}
+    return StratNodeOpScens(_strat_per(n), _branch(n), n.mult_sp, probability_branch(n), opscens)
+end
+
+Base.length(ops::StratNodeOpScens) = length(ops.opscens)
+Base.eltype(_::StratNodeOpScens) = StratNodeOperationalScenario
+
+"""
+    opscenarios(sp::StratNode{S,T,OP})
+
+Iterator that iterates over operational scenarios for a specific strategic node in the tree.
+"""
+function opscenarios(sper::StratNode{S,T,OP}) where {S,T,OP<:TimeStructure{T}}
+    return StratNodeOpScens(sper, opscenarios(sper.operational))
+end
+
+function Base.iterate(ops::StratNodeOpScens, state = (nothing, 1))
+    next =
+        isnothing(state[1]) ? iterate(ops.opscens) :
+        iterate(ops.opscens, state[1])
+    isnothing(next) && return nothing
+
+    scen = state[2]
+    return StratNodeOperationalScenario(
+        ops.sp,
+        ops.branch,
+        scen,
+        ops.mult_sp,
+        mult_scen(next[1]),
+        ops.prob_branch,
+        probability(next[1]),
+        next[1],
+    ),
+    (next[2], scen + 1)
+end
diff --git a/src/twoleveltree.jl b/src/twoleveltree.jl
index c7fe3b4..5cee0cb 100644
--- a/src/twoleveltree.jl
+++ b/src/twoleveltree.jl
@@ -129,7 +129,7 @@ Base.last(scen::StrategicScenario) = last(scen.nodes)
 function Base.iterate(scs::StrategicScenario, state = nothing)
     next = isnothing(state) ? iterate(scs.nodes) : iterate(scs.nodes, state)
     isnothing(next) && return nothing
-    return next[1].strat_node, next[2]
+    return next[1], next[2]
 end
 
 """
@@ -274,3 +274,16 @@ function Base.iterate(w::WithPrev{StratTreeNodes{T,OP}}, state) where {T, OP}
     n === nothing && return n
     return (n[1].parent, n[1]), (n[1], n[2])
 end
+
+
+"""
+    opscenarios(ts::TwoLevelTree{T,OP}) where {T,OP}
+
+Returns a collection of all operational scenarios for a [`TwoLevelTree`](@ref) time structure.
+"""
+function opscenarios(ts::TwoLevelTree{T,OP}) where {T,OP}
+    return collect(
+        Iterators.flatten(opscenarios(sp) for sp in strategic_periods(ts)),
+    )
+    return opscens
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index 6fd9919..b727664 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -931,6 +931,72 @@ end
     @test dsp[ops[4]] == 5
 end
 
+@testitem "TwoLevelTree and opscenarios" begin
+    regtree = TimeStruct.regular_tree(
+        5,
+        [3, 2],
+        OperationalScenarios(3, SimpleTimes(5, 1)),
+    )
+
+    ops1 = collect(regtree)
+    ops2 = [
+        t for sp in strat_periods(regtree) for sc in opscenarios(sp) for t in sc
+    ]
+    @test length(ops1) == length(ops2)
+    for (i, op) in enumerate(ops1)
+        @test op == ops2[i]
+    end
+
+    @test sum(length(opscenarios(sp)) for sp in strat_periods(regtree)) == 30
+    @test sum(
+        length(sc) for sp in strat_periods(regtree) for sc in opscenarios(sp)
+    ) == 150
+
+    sregtree = TimeStruct.regular_tree(5, [3, 2], SimpleTimes(5, 1))
+    ops1 = collect(sregtree)
+    ops2 = [
+        t for sp in strat_periods(sregtree) for sc in opscenarios(sp) for
+        t in sc
+    ]
+    @test length(ops1) == length(ops2)
+    @test ops1 == ops2
+
+    @test sum(length(opscenarios(sp)) for sp in strat_periods(sregtree)) == 10
+    @test sum(
+        length(sc) for sp in strat_periods(sregtree) for sc in opscenarios(sp)
+    ) == 50
+end
+
+@testitem "Strategic scenarios with operational scenarios" begin
+    regtree = TimeStruct.regular_tree(
+        5,
+        [3, 2],
+        OperationalScenarios(3, SimpleTimes(5, 1)),
+    )
+
+    @test length(TimeStruct.strategic_scenarios(regtree)) == 6
+
+    for sc in TimeStruct.strategic_scenarios(regtree)
+        @test length(sc) == length(collect(sc))
+
+        for (prev_sp, sp) in withprev(sc)
+            if !isnothing(prev_sp)
+                @test TimeStruct._strat_per(prev_sp) + 1 ==
+                      TimeStruct._strat_per(sp)
+            end
+        end
+    end
+end
+
+@testitem "TwoLevel as a tree" begin
+    two_level = TwoLevel(5, 10, SimpleTimes(10, 1))
+
+    scens = TimeStruct.strategic_scenarios(two_level)
+    @test length(scens) == 1
+    sps = collect(sp for sc in TimeStruct.strategic_scenarios(two_level) for sp in strat_periods(sc))
+    @test length(sps) == 5
+end
+
 @testitem "Profiles constructors" begin
     # Checking the input type
     @test_throws MethodError FixedProfile("wrong_input")
@@ -1026,72 +1092,6 @@ end
     @test vals == [1, 2, 3, 4, 2, 4, 6, 8, 3, 6, 9, 12]
 end
 
-@testitem "TwoLevelTree and opscenarios" begin
-    regtree = TimeStruct.regular_tree(
-        5,
-        [3, 2],
-        OperationalScenarios(3, SimpleTimes(5, 1)),
-    )
-
-    ops1 = collect(regtree)
-    ops2 = [
-        t for sp in strat_periods(regtree) for sc in opscenarios(sp) for t in sc
-    ]
-    @test length(ops1) == length(ops2)
-    for (i, op) in enumerate(ops1)
-        @test op == ops2[i]
-    end
-
-    @test sum(length(opscenarios(sp)) for sp in strat_periods(regtree)) == 30
-    @test sum(
-        length(sc) for sp in strat_periods(regtree) for sc in opscenarios(sp)
-    ) == 150
-
-    sregtree = TimeStruct.regular_tree(5, [3, 2], SimpleTimes(5, 1))
-    ops1 = collect(sregtree)
-    ops2 = [
-        t for sp in strat_periods(sregtree) for sc in opscenarios(sp) for
-        t in sc
-    ]
-    @test length(ops1) == length(ops2)
-    @test ops1 == ops2
-
-    @test sum(length(opscenarios(sp)) for sp in strat_periods(sregtree)) == 10
-    @test sum(
-        length(sc) for sp in strat_periods(sregtree) for sc in opscenarios(sp)
-    ) == 50
-end
-
-@testitem "Strategic scenarios with operational scenarios" begin
-    regtree = TimeStruct.regular_tree(
-        5,
-        [3, 2],
-        OperationalScenarios(3, SimpleTimes(5, 1)),
-    )
-
-    @test length(strategic_scenarios(regtree)) == 6
-
-    for sc in strategic_scenarios(regtree)
-        @test length(sc) == length(collect(sc))
-
-        for (prev_sp, sp) in withprev(sc)
-            if !isnothing(prev_sp)
-                @test TimeStruct._strat_per(prev_sp) + 1 ==
-                      TimeStruct._strat_per(sp)
-            end
-        end
-    end
-end
-
-@testitem "TwoLevel as a tree" begin
-    two_level = TwoLevel(5, 10, SimpleTimes(10, 1))
-
-    scens = TimeStruct.strategic_scenarios(two_level)
-    @test length(scens) == 1
-    sps = collect(sp for sc in TimeStruct.strategic_scenarios(two_level) for sp in strat_periods(sc))
-    @test length(sps) == 5
-end
-
 @testitem "Iteration utilities" begin
     uniform_day = SimpleTimes(24, 1)
     uniform_week = TwoLevel(7, 24, uniform_day)

From c231baafce95d13f7e46b64f5b1b1a3ee40b80dc Mon Sep 17 00:00:00 2001
From: Julian Straus <julian.straus@sintef.no>
Date: Tue, 27 Aug 2024 15:05:26 +0200
Subject: [PATCH 04/14] Added support for RepresentativePeriods

---
 src/tree_nodes.jl   | 114 ++++++++++++++++++++++++++++++++++++++++++--
 src/twoleveltree.jl |  11 +++++
 test/runtests.jl    |  36 ++++++++++++++
 3 files changed, 157 insertions(+), 4 deletions(-)

diff --git a/src/tree_nodes.jl b/src/tree_nodes.jl
index 60c8fc6..b6c4ea4 100644
--- a/src/tree_nodes.jl
+++ b/src/tree_nodes.jl
@@ -109,8 +109,6 @@ function Base.iterate(os::StratNodeOperationalScenario, state = nothing)
     next[2]
 end
 
-
-# Iteration through scenarios
 """
     struct StratNodeOpScens
 
@@ -137,8 +135,8 @@ Base.eltype(_::StratNodeOpScens) = StratNodeOperationalScenario
 
 Iterator that iterates over operational scenarios for a specific strategic node in the tree.
 """
-function opscenarios(sper::StratNode{S,T,OP}) where {S,T,OP<:TimeStructure{T}}
-    return StratNodeOpScens(sper, opscenarios(sper.operational))
+function opscenarios(sn::StratNode{S,T,OP}) where {S,T,OP<:TimeStructure{T}}
+    return StratNodeOpScens(sn, opscenarios(sn.operational))
 end
 
 function Base.iterate(ops::StratNodeOpScens, state = (nothing, 1))
@@ -160,3 +158,111 @@ function Base.iterate(ops::StratNodeOpScens, state = (nothing, 1))
     ),
     (next[2], scen + 1)
 end
+
+
+"""
+    struct StratNodeReprPeriod{T,OP<:TimeStructure{T}} <: AbstractRepresentativePeriod{T}
+
+A structure representing a single representative period of a strategic node supporting
+iteration over its time periods.
+"""
+struct StratNodeReprPeriod{T,OP<:TimeStructure{T}} <: AbstractRepresentativePeriod{T}
+    sp::Int
+    branch::Int
+    rp::Int
+    mult_sp::Float64
+    mult_rp::Float64
+    prob_branch::Float64
+    operational::OP
+end
+
+function Base.show(io::IO, srp::StratNodeReprPeriod)
+    return print(io, "sp$(srp.sp)-br$(srp.branch)-sc$(srp.rp)")
+end
+
+probability(srp::StratNodeReprPeriod) = srp.prob_branch
+multiple(srp::StratNodeReprPeriod, t::OperationalPeriod) = t.multiple / srp.mult_sp
+_rper(srp::StratNodeReprPeriod) = srp.rp
+_branch(srp::StratNodeReprPeriod) = srp.branch
+_strat_per(srp::StratNodeReprPeriod) = srp.sp
+
+StrategicTreeIndexable(::Type{<:StratNodeReprPeriod}) = HasStratTreeIndex()
+StrategicIndexable(::Type{<:StratNodeReprPeriod}) = HasStratIndex()
+
+Base.length(snrp::StratNodeReprPeriod) = length(snrp.operational)
+Base.eltype(_::StratNodeReprPeriod) = TreePeriod
+
+function Base.last(srp::StratNodeReprPeriod)
+    per = last(srp.operational)
+    return OperationalPeriod(srp.sp, per, srp.mult_sp * multiple(per))
+end
+
+function Base.getindex(srp::StratNodeReprPeriod, index)
+    per = srp.operational[index]
+    mult = srp.mult_sp * multiple(per)
+    return TreePeriod(srp.sp, srp.branch, probability_branch(srp), mult, per)
+end
+
+function Base.eachindex(srp::StratNodeReprPeriod)
+    return eachindex(srp.operational)
+end
+
+# Iterate the time periods of a StratOperationalScenario
+function Base.iterate(srp::StratNodeReprPeriod, state = nothing)
+    next =
+        isnothing(state) ? iterate(srp.operational) :
+        iterate(srp.operational, state)
+    isnothing(next) && return nothing
+
+    return TreePeriod(srp.sp, srp.branch, srp.prob_branch, srp.mult_sp * multiple(next[1]), next[1]),
+    next[2]
+end
+
+"""
+    struct StratNodeReprPeriods
+
+Type for iterating through the individual operational scenarios of a [`StratNode`](@ref).
+It is automatically created through the function [`repr_periods`](@ref).
+"""
+struct StratNodeReprPeriods
+    sp::Int
+    branch::Int
+    mult_sp::Float64
+    prob_branch::Float64
+    repr::Any
+end
+
+function StratNodeReprPeriods(n::StratNode{S,T,OP}, repr) where {S,T,OP<:TimeStructure{T}}
+    return StratNodeReprPeriods(_strat_per(n), _branch(n), n.mult_sp, probability_branch(n), repr)
+end
+
+Base.length(ops::StratNodeReprPeriods) = length(ops.repr)
+Base.eltype(_::StratNodeReprPeriods) = StratNodeReprPeriod
+
+"""
+    repr_periods(sp::StratNode{S,T,OP})
+
+Iterator that iterates over operational scenarios for a specific strategic node in the tree.
+"""
+function repr_periods(sn::StratNode{S,T,OP}) where {S,T,OP<:TimeStructure{T}}
+    return StratNodeReprPeriods(sn, repr_periods(sn.operational))
+end
+
+function Base.iterate(ops::StratNodeReprPeriods, state = (nothing, 1))
+    next =
+        isnothing(state[1]) ? iterate(ops.repr) :
+        iterate(ops.repr, state[1])
+    isnothing(next) && return nothing
+
+    scen = state[2]
+    return StratNodeReprPeriod(
+        ops.sp,
+        ops.branch,
+        scen,
+        ops.mult_sp,
+        mult_repr(next[1]),
+        ops.prob_branch,
+        next[1],
+    ),
+    (next[2], scen + 1)
+end
diff --git a/src/twoleveltree.jl b/src/twoleveltree.jl
index 5cee0cb..1d22335 100644
--- a/src/twoleveltree.jl
+++ b/src/twoleveltree.jl
@@ -287,3 +287,14 @@ function opscenarios(ts::TwoLevelTree{T,OP}) where {T,OP}
     )
     return opscens
 end
+
+"""
+    repr_periods(ts::TwoLevelTree{T,OP})
+
+Returns a collection of all representative periods for a [`TwoLevelTree`](@ref) time structure.
+"""
+function repr_periods(ts::TwoLevelTree{T,OP}) where {T,OP}
+    return collect(
+        Iterators.flatten(repr_periods(sp) for sp in strategic_periods(ts)),
+    )
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index b727664..a3603d2 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -967,6 +967,42 @@ end
     ) == 50
 end
 
+@testitem "TwoLevelTree and repr_periods" begin
+    regtree = TimeStruct.regular_tree(
+        5,
+        [3, 2],
+        RepresentativePeriods(2, 1, SimpleTimes(5, 1)),
+    )
+
+    ops1 = collect(regtree)
+    ops2 = [
+        t for sp in strat_periods(regtree) for sc in repr_periods(sp) for t in sc
+    ]
+    @test length(ops1) == length(ops2)
+    for (i, op) in enumerate(ops1)
+        @test op == ops2[i]
+    end
+
+    @test sum(length(repr_periods(sp)) for sp in strat_periods(regtree)) == 20
+    @test sum(
+        length(rper) for sp in strat_periods(regtree) for rper in repr_periods(sp)
+    ) == 100
+
+    sregtree = TimeStruct.regular_tree(5, [3, 2], SimpleTimes(5, 1))
+    ops1 = collect(sregtree)
+    ops2 = [
+        t for sp in strat_periods(sregtree) for rper in repr_periods(sp) for
+        t in rper
+    ]
+    @test length(ops1) == length(ops2)
+    @test ops1 == ops2
+
+    @test sum(length(repr_periods(sp)) for sp in strat_periods(sregtree)) == 10
+    @test sum(
+        length(rper) for sp in strat_periods(sregtree) for rper in repr_periods(sp)
+    ) == 50
+end
+
 @testitem "Strategic scenarios with operational scenarios" begin
     regtree = TimeStruct.regular_tree(
         5,

From 541d75b4c215a1869a60d75dbc25c7f5e1321f59 Mon Sep 17 00:00:00 2001
From: Julian Straus <julian.straus@sintef.no>
Date: Wed, 28 Aug 2024 08:35:22 +0200
Subject: [PATCH 05/14] Added support for RepresentativePeriods with
 OperationalScenarios

---
 src/tree_nodes.jl | 139 ++++++++++++++++++++++++++++++++++++++++++++++
 test/runtests.jl  |  47 ++++++++++++++++
 2 files changed, 186 insertions(+)

diff --git a/src/tree_nodes.jl b/src/tree_nodes.jl
index b6c4ea4..c3ea6fd 100644
--- a/src/tree_nodes.jl
+++ b/src/tree_nodes.jl
@@ -181,6 +181,7 @@ function Base.show(io::IO, srp::StratNodeReprPeriod)
 end
 
 probability(srp::StratNodeReprPeriod) = srp.prob_branch
+probability_branch(srp::StratNodeReprPeriod) = srp.prob_branch
 multiple(srp::StratNodeReprPeriod, t::OperationalPeriod) = t.multiple / srp.mult_sp
 _rper(srp::StratNodeReprPeriod) = srp.rp
 _branch(srp::StratNodeReprPeriod) = srp.branch
@@ -266,3 +267,141 @@ function Base.iterate(ops::StratNodeReprPeriods, state = (nothing, 1))
     ),
     (next[2], scen + 1)
 end
+
+
+"""
+    struct StratNodeReprOpscenario{T} <: AbstractOperationalScenario{T}
+
+A structure representing a single operational scenario for a representative period in A
+[`TwoLevelTree`](@ref) structure supporting iteration over its time periods.
+"""
+struct StratNodeReprOpscenario{T} <: AbstractOperationalScenario{T}
+    sp::Int
+    branch::Int
+    rp::Int
+    opscen::Int
+    mult_sp::Float64
+    mult_rp::Float64
+    prob_branch::Float64
+    prob_scen::Float64
+    operational::TimeStructure{T}
+end
+
+probability(os::StratNodeReprOpscenario) = os.prob_branch * os.prob_scen
+probability_branch(os::StratNodeReprOpscenario) = os.prob_branch
+_opscen(os::StratNodeReprOpscenario) = os.opscen
+_rper(os::StratNodeReprOpscenario) = os.rp
+_branch(os::StratNodeReprOpscenario) = os.branch
+_strat_per(os::StratNodeReprOpscenario) = os.sp
+
+function Base.show(io::IO, srop::StratNodeReprOpscenario)
+    return print(io, "sp$(srop.sp)-br$(os.branch)-rp$(srop.rp)-sc$(srop.opscen)")
+end
+
+StrategicTreeIndexable(::Type{<:StratNodeReprOpscenario}) = HasStratTreeIndex()
+StrategicIndexable(::Type{<:StratNodeReprOpscenario}) = HasStratIndex()
+function RepresentativeIndexable(::Type{<:StratNodeReprOpscenario})
+    return HasReprIndex()
+end
+ScenarioIndexable(::Type{<:StratNodeReprOpscenario}) = HasScenarioIndex()
+
+function Base.last(os::StratNodeReprOpscenario)
+    per = last(os.operational)
+    rper = ReprPeriod(os.rp, per, os.mult_rp * multiple(per))
+    return TreePeriod(
+        os.sp,
+        os.branch,
+        probability_branch(os),
+        os.mult_sp * os.mult_rp * multiple(per),
+        rper,
+    )
+end
+
+function Base.getindex(os::StratNodeReprOpscenario, index)
+    mult = stripunit(os.duration * os.op_per_strat / duration(os.operational))
+    period = ReprPeriod(os.rp, os.operational[index], mult)
+    return TreePeriod(os.sp, os.branch, probability_branch(os), mult, period)
+end
+
+function Base.eachindex(os::StratNodeReprOpscenario)
+    return eachindex(os.operational)
+end
+
+# Iterate the time periods of a StratNodeReprOpscenario
+function Base.iterate(os::StratNodeReprOpscenario, state = nothing)
+    next =
+        isnothing(state) ? iterate(os.operational) :
+        iterate(os.operational, state)
+    isnothing(next) && return nothing
+
+    period = ReprPeriod(os.rp, next[1], os.mult_rp * multiple(next[1]))
+    return TreePeriod(
+        os.sp,
+        os.branch,
+        probability_branch(os),
+        os.mult_sp * os.mult_rp * multiple(next[1]),
+        period,
+    ),
+    next[2]
+end
+
+Base.length(os::StratNodeReprOpscenario) = length(os.operational)
+
+function Base.eltype(::Type{StratNodeReprOpscenario{T}}) where {T}
+    return TreePeriod
+end
+
+"""
+    struct StratNodeReprOpscenarios
+
+Type for iterating through the individual operational scenarios of a
+[`StratNodeReprPeriod`](@ref). It is automatically created through the function
+[`opscenarios`](@ref).
+"""
+struct StratNodeReprOpscenarios
+    sp::Int
+    branch::Int
+    rp::Int
+    mult_sp::Float64
+    mult_rp::Float64
+    prob_branch::Float64
+    opscens::Any
+end
+
+function StratNodeReprOpscenarios(n::StratNodeReprPeriod{T,OP}, opscens) where {T,OP<:TimeStructure{T}}
+    return StratNodeReprOpscenarios(_strat_per(n), _branch(n), _rper(n), n.mult_sp, n.mult_rp, probability_branch(n), opscens)
+end
+
+"""
+    opscenarios(sp::StratNodeReprPeriod{T,RepresentativePeriod{T,OP}})
+
+Iterator that iterates over operational scenarios for a specific strategic node in the tree.
+"""
+function opscenarios(
+    srp::StratNodeReprPeriod{T,RepresentativePeriod{T,OP}},
+) where {T,OP}
+    return StratNodeReprOpscenarios(srp, opscenarios(srp.operational.operational))
+end
+
+Base.length(srop::StratNodeReprOpscenarios) = length(srop.opscens)
+
+function Base.iterate(srop::StratNodeReprOpscenarios, state = (nothing, 1))
+    next =
+        isnothing(state[1]) ? iterate(srop.opscens) :
+        iterate(srop.opscens, state[1])
+    isnothing(next) && return nothing
+
+    scen = state[2]
+    return StratNodeReprOpscenario(
+        srop.sp,
+        srop.branch,
+        srop.rp,
+        scen,
+        srop.mult_sp,
+        srop.mult_rp,
+        srop.prob_branch,
+        probability(next[1]),
+        next[1],
+    ),
+    (next[2], scen + 1)
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index a3603d2..156215a 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1003,6 +1003,53 @@ end
     ) == 50
 end
 
+@testitem "TwoLevelTree with RepresentativePeriods and OperationalScenarios" begin
+    regtree = TimeStruct.regular_tree(
+        5,
+        [3, 2],
+        RepresentativePeriods(2, 1, OperationalScenarios(3, SimpleTimes(5, 1))),
+    )
+
+    ops1 = collect(regtree)
+    ops2 = [
+        t for sp in strat_periods(regtree) for t in sp
+    ]
+    ops3 = [
+        t for sp in strat_periods(regtree) for sc in opscenarios(sp) for t in sc
+    ]
+    ops4 = [
+        t for sp in strat_periods(regtree) for sc in repr_periods(sp) for t in sc
+    ]
+    ops5 = [
+        t for sp in strat_periods(regtree) for rp in repr_periods(sp) for sc in opscenarios(rp) for t in sc
+    ]
+    @test length(ops1) == length(ops2)
+    @test length(ops1) == length(ops3)
+    @test length(ops1) == length(ops4)
+    @test length(ops1) == length(ops5)
+    for (i, op) in enumerate(ops1)
+        @test op == ops2[i]
+        @test op == ops3[i]
+        @test op == ops4[i]
+        @test op == ops5[i]
+    end
+
+    @test sum(length(repr_periods(sp)) for sp in strat_periods(regtree)) == 20
+    @test sum(length(opscenarios(sp)) for sp in strat_periods(regtree)) == 60
+    @test sum(
+        length(rp) for sp in strat_periods(regtree) for rp in repr_periods(sp)
+    ) == 300
+    @test sum(
+        length(rp) for rp in repr_periods(regtree)
+    ) == 300
+    @test sum(
+        length(sc) for sp in strat_periods(regtree) for sc in opscenarios(sp)
+    ) == 300
+    @test sum(
+        length(sc) for sp in strat_periods(regtree) for rp in repr_periods(sp) for sc in opscenarios(rp)
+    ) == 300
+end
+
 @testitem "Strategic scenarios with operational scenarios" begin
     regtree = TimeStruct.regular_tree(
         5,

From e2d967e91eb8e739f22756f3953edf9de5f62bab Mon Sep 17 00:00:00 2001
From: Julian Straus <julian.straus@sintef.no>
Date: Wed, 28 Aug 2024 10:12:00 +0200
Subject: [PATCH 06/14] Unified the structure of individual types

---
 src/tree_nodes.jl   | 148 ++++++++++++++++++----------------
 src/twoleveltree.jl |  35 +++++++--
 test/Project.toml   |   3 +-
 test/runtests.jl    | 188 +++++++++++++++++++-------------------------
 4 files changed, 191 insertions(+), 183 deletions(-)

diff --git a/src/tree_nodes.jl b/src/tree_nodes.jl
index c3ea6fd..6cbb151 100644
--- a/src/tree_nodes.jl
+++ b/src/tree_nodes.jl
@@ -5,6 +5,24 @@ Abstract base type for all tree nodes within a [`TwoLevelTree`] type
 """
 abstract type AbstractTreeNode{T} <: TimeStructure{T} end
 
+abstract type StrategicTreeIndexable end
+struct HasStratTreeIndex <: StrategicTreeIndexable end
+struct NoStratTreeIndex <: StrategicTreeIndexable end
+
+StrategicTreeIndexable(::Type) = NoStratTreeIndex()
+StrategicTreeIndexable(::Type{<:AbstractTreeNode}) = HasStratTreeIndex()
+StrategicTreeIndexable(::Type{<:TimePeriod}) = HasStratTreeIndex()
+StrategicIndexable(::Type{<:AbstractTreeNode}) = HasStratIndex()
+
+"""
+    struct StratNode{S, T, OP<:TimeStructure{T}} <: AbstractTreeNode{T}
+
+A structure representing a single strategic node of a [`TwolevelTree`](@ref). It is created
+through iterating through [`StratTreeNodes`](@ref).
+
+It is equivalent to a [`StrategicPeriod`](@ref) of a [`TwoLevel`](@ref) time structure when
+utilizing a [`TwolevelTree`](@ref).
+"""
 struct StratNode{S, T, OP<:TimeStructure{T}} <: AbstractTreeNode{T}
     sp::Int
     branch::Int
@@ -15,18 +33,19 @@ struct StratNode{S, T, OP<:TimeStructure{T}} <: AbstractTreeNode{T}
     operational::OP
 end
 
-Base.show(io::IO, n::StratNode) = print(io, "sp$(n.sp)-br$(n.branch)")
-_branch(n::StratNode) = n.branch
 _strat_per(n::StratNode) = n.sp
+_branch(n::StratNode) = n.branch
+
 probability_branch(n::StratNode) = n.prob_branch
 duration_strat(n::StratNode) = n.duration
 multiple_strat(sp::StratNode, t) = multiple(t) / duration_strat(sp)
 
 isfirst(n::StratNode) = n.sp == 1
 
-# Iterate through time periods of a strategic node
+# Adding methods to existing Julia functions
+Base.show(io::IO, n::StratNode) = print(io, "sp$(n.sp)-br$(n.branch)")
 Base.length(n::StratNode) = length(n.operational)
-Base.eltype(::Type{StratNode{T}}) where {T} = TreePeriod
+Base.eltype(::Type{StratNode}) = TreePeriod
 function Base.iterate(itr::StratNode, state = nothing)
     next =
         isnothing(state) ? iterate(itr.operational) :
@@ -39,22 +58,15 @@ function Base.iterate(itr::StratNode, state = nothing)
     next[2]
 end
 
-abstract type StrategicTreeIndexable end
-
-struct HasStratTreeIndex <: StrategicTreeIndexable end
-struct NoStratTreeIndex <: StrategicTreeIndexable end
-
-StrategicTreeIndexable(::Type) = NoStratTreeIndex()
-StrategicTreeIndexable(::Type{<:AbstractTreeNode}) = HasStratTreeIndex()
-StrategicTreeIndexable(::Type{<:TimePeriod}) = HasStratTreeIndex()
-StrategicIndexable(::Type{<:AbstractTreeNode}) = HasStratIndex()
-
-
 """
     struct StratNodeOperationalScenario{T} <: AbstractOperationalScenario{T}
 
 A structure representing a single operational scenario for a strategic node supporting
-iteration over its time periods.
+iteration over its time periods. It is created through iterating through
+[`StratNodeOpScens`](@ref).
+
+It is equivalent to a [`StratOperationalScenario`](@ref) of a [`TwoLevel`](@ref) time
+structure when utilizing a [`TwolevelTree`](@ref).
 """
 struct StratNodeOperationalScenario{T} <: AbstractOperationalScenario{T}
     sp::Int
@@ -67,45 +79,44 @@ struct StratNodeOperationalScenario{T} <: AbstractOperationalScenario{T}
     operational::TimeStructure{T}
 end
 
-function Base.show(io::IO, os::StratNodeOperationalScenario)
-    return print(io, "sp$(os.sp)-br$(os.branch)-sc$(os.scen)")
-end
-
-probability(os::StratNodeOperationalScenario) = os.prob_branch * prob_scen
-mult_scen(os::StratNodeOperationalScenario) = os.mult_scen
 _opscen(os::StratNodeOperationalScenario) = os.scen
 _branch(os::StratNodeOperationalScenario) = os.branch
 _strat_per(os::StratNodeOperationalScenario) = os.sp
 
+probability(os::StratNodeOperationalScenario) = os.prob_branch * prob_scen
+probability_branch(os::StratNodeOperationalScenario) = os.prob_branch
+mult_scen(os::StratNodeOperationalScenario) = os.mult_scen
+
 StrategicTreeIndexable(::Type{<:StratNodeOperationalScenario}) = HasStratTreeIndex()
 StrategicIndexable(::Type{<:StratNodeOperationalScenario}) = HasStratIndex()
 
+# Adding methods to existing Julia functions
+function Base.show(io::IO, os::StratNodeOperationalScenario)
+    return print(io, "sp$(os.sp)-br$(os.branch)-sc$(os.scen)")
+end
 Base.length(snops::StratNodeOperationalScenario) = length(snops.operational)
 Base.eltype(_::StratNodeOperationalScenario) = TreePeriod
-
 function Base.last(os::StratNodeOperationalScenario)
     per = last(os.operational)
-    return OperationalPeriod(os.sp, per, os.mult_sp * multiple(per))
+    mult = os.mult_sp * multiple(per)
+    return TreePeriod(os.sp, os.branch, probability_branch(os), mult, per)
 end
-
 function Base.getindex(os::StratNodeOperationalScenario, index)
     per = os.operational[index]
     mult = os.mult_sp * multiple(per)
     return TreePeriod(os.sp, os.branch, probability_branch(os), mult, per)
 end
-
 function Base.eachindex(os::StratNodeOperationalScenario)
     return eachindex(os.operational)
 end
-
-# Iterate the time periods of a StratOperationalScenario
 function Base.iterate(os::StratNodeOperationalScenario, state = nothing)
     next =
         isnothing(state) ? iterate(os.operational) :
         iterate(os.operational, state)
     isnothing(next) && return nothing
 
-    return TreePeriod(os.sp, os.branch, os.prob_branch, os.mult_sp * multiple(next[1]), next[1]),
+    mult = os.mult_sp * multiple(next[1])
+    return TreePeriod(os.sp, os.branch, os.prob_branch, mult, next[1]),
     next[2]
 end
 
@@ -127,18 +138,18 @@ function StratNodeOpScens(n::StratNode{S,T,OP}, opscens) where {S,T,OP<:TimeStru
     return StratNodeOpScens(_strat_per(n), _branch(n), n.mult_sp, probability_branch(n), opscens)
 end
 
-Base.length(ops::StratNodeOpScens) = length(ops.opscens)
-Base.eltype(_::StratNodeOpScens) = StratNodeOperationalScenario
-
 """
     opscenarios(sp::StratNode{S,T,OP})
 
-Iterator that iterates over operational scenarios for a specific strategic node in the tree.
+When the `TimeStructure` is a [`StratNode`](@ref), `opscenarios` returns a
+[`StratNodeOpScens`](@ref) used for iterating through the individual operational scenarios
 """
 function opscenarios(sn::StratNode{S,T,OP}) where {S,T,OP<:TimeStructure{T}}
     return StratNodeOpScens(sn, opscenarios(sn.operational))
 end
 
+Base.length(ops::StratNodeOpScens) = length(ops.opscens)
+Base.eltype(_::StratNodeOpScens) = StratNodeOperationalScenario
 function Base.iterate(ops::StratNodeOpScens, state = (nothing, 1))
     next =
         isnothing(state[1]) ? iterate(ops.opscens) :
@@ -163,8 +174,11 @@ end
 """
     struct StratNodeReprPeriod{T,OP<:TimeStructure{T}} <: AbstractRepresentativePeriod{T}
 
-A structure representing a single representative period of a strategic node supporting
-iteration over its time periods.
+A structure representing a single representative period of a [`StrategicNode`](@ref) of a
+[`TwolevelTree`](@ref). It is created through iterating through [`StratNodeReprPeriods`](@ref).
+
+It is equivalent to a [`StratReprPeriod`](@ref) of a [`TwoLevel`](@ref) time structure when
+utilizing a [`TwolevelTree`](@ref).
 """
 struct StratNodeReprPeriod{T,OP<:TimeStructure{T}} <: AbstractRepresentativePeriod{T}
     sp::Int
@@ -176,53 +190,52 @@ struct StratNodeReprPeriod{T,OP<:TimeStructure{T}} <: AbstractRepresentativePeri
     operational::OP
 end
 
-function Base.show(io::IO, srp::StratNodeReprPeriod)
-    return print(io, "sp$(srp.sp)-br$(srp.branch)-sc$(srp.rp)")
-end
 
-probability(srp::StratNodeReprPeriod) = srp.prob_branch
-probability_branch(srp::StratNodeReprPeriod) = srp.prob_branch
-multiple(srp::StratNodeReprPeriod, t::OperationalPeriod) = t.multiple / srp.mult_sp
 _rper(srp::StratNodeReprPeriod) = srp.rp
 _branch(srp::StratNodeReprPeriod) = srp.branch
 _strat_per(srp::StratNodeReprPeriod) = srp.sp
 
+probability(srp::StratNodeReprPeriod) = srp.prob_branch
+probability_branch(srp::StratNodeReprPeriod) = srp.prob_branch
+multiple(srp::StratNodeReprPeriod, t::OperationalPeriod) = t.multiple / srp.mult_sp
+
 StrategicTreeIndexable(::Type{<:StratNodeReprPeriod}) = HasStratTreeIndex()
 StrategicIndexable(::Type{<:StratNodeReprPeriod}) = HasStratIndex()
 
+# Adding methods to existing Julia functions
+function Base.show(io::IO, srp::StratNodeReprPeriod)
+    return print(io, "sp$(srp.sp)-br$(srp.branch)-sc$(srp.rp)")
+end
 Base.length(snrp::StratNodeReprPeriod) = length(snrp.operational)
 Base.eltype(_::StratNodeReprPeriod) = TreePeriod
-
 function Base.last(srp::StratNodeReprPeriod)
     per = last(srp.operational)
-    return OperationalPeriod(srp.sp, per, srp.mult_sp * multiple(per))
+    mult = srp.mult_sp * multiple(per)
+    return TreePeriod(srp.sp, srp.branch, probability_branch(srp), mult, per)
 end
-
 function Base.getindex(srp::StratNodeReprPeriod, index)
     per = srp.operational[index]
     mult = srp.mult_sp * multiple(per)
     return TreePeriod(srp.sp, srp.branch, probability_branch(srp), mult, per)
 end
-
 function Base.eachindex(srp::StratNodeReprPeriod)
     return eachindex(srp.operational)
 end
-
-# Iterate the time periods of a StratOperationalScenario
 function Base.iterate(srp::StratNodeReprPeriod, state = nothing)
     next =
         isnothing(state) ? iterate(srp.operational) :
         iterate(srp.operational, state)
     isnothing(next) && return nothing
 
-    return TreePeriod(srp.sp, srp.branch, srp.prob_branch, srp.mult_sp * multiple(next[1]), next[1]),
+    mult = srp.mult_sp * multiple(next[1])
+    return TreePeriod(srp.sp, srp.branch, srp.prob_branch, mult, next[1]),
     next[2]
 end
 
 """
     struct StratNodeReprPeriods
 
-Type for iterating through the individual operational scenarios of a [`StratNode`](@ref).
+Type for iterating through the individual presentative periods of a [`StratNode`](@ref).
 It is automatically created through the function [`repr_periods`](@ref).
 """
 struct StratNodeReprPeriods
@@ -237,9 +250,6 @@ function StratNodeReprPeriods(n::StratNode{S,T,OP}, repr) where {S,T,OP<:TimeStr
     return StratNodeReprPeriods(_strat_per(n), _branch(n), n.mult_sp, probability_branch(n), repr)
 end
 
-Base.length(ops::StratNodeReprPeriods) = length(ops.repr)
-Base.eltype(_::StratNodeReprPeriods) = StratNodeReprPeriod
-
 """
     repr_periods(sp::StratNode{S,T,OP})
 
@@ -249,6 +259,8 @@ function repr_periods(sn::StratNode{S,T,OP}) where {S,T,OP<:TimeStructure{T}}
     return StratNodeReprPeriods(sn, repr_periods(sn.operational))
 end
 
+Base.length(ops::StratNodeReprPeriods) = length(ops.repr)
+Base.eltype(_::StratNodeReprPeriods) = StratNodeReprPeriod
 function Base.iterate(ops::StratNodeReprPeriods, state = (nothing, 1))
     next =
         isnothing(state[1]) ? iterate(ops.repr) :
@@ -268,7 +280,6 @@ function Base.iterate(ops::StratNodeReprPeriods, state = (nothing, 1))
     (next[2], scen + 1)
 end
 
-
 """
     struct StratNodeReprOpscenario{T} <: AbstractOperationalScenario{T}
 
@@ -287,16 +298,13 @@ struct StratNodeReprOpscenario{T} <: AbstractOperationalScenario{T}
     operational::TimeStructure{T}
 end
 
-probability(os::StratNodeReprOpscenario) = os.prob_branch * os.prob_scen
-probability_branch(os::StratNodeReprOpscenario) = os.prob_branch
 _opscen(os::StratNodeReprOpscenario) = os.opscen
 _rper(os::StratNodeReprOpscenario) = os.rp
 _branch(os::StratNodeReprOpscenario) = os.branch
 _strat_per(os::StratNodeReprOpscenario) = os.sp
 
-function Base.show(io::IO, srop::StratNodeReprOpscenario)
-    return print(io, "sp$(srop.sp)-br$(os.branch)-rp$(srop.rp)-sc$(srop.opscen)")
-end
+probability(os::StratNodeReprOpscenario) = os.prob_branch * os.prob_scen
+probability_branch(os::StratNodeReprOpscenario) = os.prob_branch
 
 StrategicTreeIndexable(::Type{<:StratNodeReprOpscenario}) = HasStratTreeIndex()
 StrategicIndexable(::Type{<:StratNodeReprOpscenario}) = HasStratIndex()
@@ -305,22 +313,29 @@ function RepresentativeIndexable(::Type{<:StratNodeReprOpscenario})
 end
 ScenarioIndexable(::Type{<:StratNodeReprOpscenario}) = HasScenarioIndex()
 
+# Adding methods to existing Julia functions
+function Base.show(io::IO, srop::StratNodeReprOpscenario)
+    return print(io, "sp$(srop.sp)-br$(os.branch)-rp$(srop.rp)-sc$(srop.opscen)")
+end
+Base.length(srop::StratNodeReprOpscenario) = length(srop.operational)
+Base.eltype(_::StratNodeReprOpscenario) = TreePeriod
 function Base.last(os::StratNodeReprOpscenario)
     per = last(os.operational)
     rper = ReprPeriod(os.rp, per, os.mult_rp * multiple(per))
+    mult - os.mult_sp * os.mult_rp * multiple(per)
     return TreePeriod(
         os.sp,
         os.branch,
         probability_branch(os),
-        os.mult_sp * os.mult_rp * multiple(per),
+        mult,
         rper,
     )
 end
 
 function Base.getindex(os::StratNodeReprOpscenario, index)
     mult = stripunit(os.duration * os.op_per_strat / duration(os.operational))
-    period = ReprPeriod(os.rp, os.operational[index], mult)
-    return TreePeriod(os.sp, os.branch, probability_branch(os), mult, period)
+    per = ReprPeriod(os.rp, os.operational[index], mult)
+    return TreePeriod(os.sp, os.branch, probability_branch(os), mult, per)
 end
 
 function Base.eachindex(os::StratNodeReprOpscenario)
@@ -335,22 +350,17 @@ function Base.iterate(os::StratNodeReprOpscenario, state = nothing)
     isnothing(next) && return nothing
 
     period = ReprPeriod(os.rp, next[1], os.mult_rp * multiple(next[1]))
+    mult = os.mult_sp * os.mult_rp * multiple(next[1])
     return TreePeriod(
         os.sp,
         os.branch,
         probability_branch(os),
-        os.mult_sp * os.mult_rp * multiple(next[1]),
+        mult,
         period,
     ),
     next[2]
 end
 
-Base.length(os::StratNodeReprOpscenario) = length(os.operational)
-
-function Base.eltype(::Type{StratNodeReprOpscenario{T}}) where {T}
-    return TreePeriod
-end
-
 """
     struct StratNodeReprOpscenarios
 
@@ -384,7 +394,7 @@ function opscenarios(
 end
 
 Base.length(srop::StratNodeReprOpscenarios) = length(srop.opscens)
-
+Base.eltype(_::StratNodeReprOpscenarios) = StratNodeReprOpscenario
 function Base.iterate(srop::StratNodeReprOpscenarios, state = (nothing, 1))
     next =
         isnothing(state[1]) ? iterate(srop.opscens) :
diff --git a/src/twoleveltree.jl b/src/twoleveltree.jl
index 1d22335..9dba329 100644
--- a/src/twoleveltree.jl
+++ b/src/twoleveltree.jl
@@ -248,9 +248,10 @@ although the meaining remains unchanged.
 struct StratTreeNodes{T, OP}
     ts::TwoLevelTree{T, OP}
 end
-strat_periods(ts::TwoLevelTree) = StratTreeNodes(ts)
-Base.length(sps::StratTreeNodes) = length(sps.ts.nodes)
 
+# Adding methods to existing Julia functions
+Base.length(sps::StratTreeNodes) = length(sps.ts.nodes)
+Base.eltype(_::Type{StratTreeNodes{T, OP}}) where{T, OP} = OP
 function Base.iterate(stps::StratTreeNodes, state = nothing)
     next = isnothing(state) ? 1 : state + 1
     next == length(stps) + 1 && return nothing
@@ -275,13 +276,28 @@ function Base.iterate(w::WithPrev{StratTreeNodes{T,OP}}, state) where {T, OP}
     return (n[1].parent, n[1]), (n[1], n[2])
 end
 
+"""
+    strat_periods(ts::TwoLevelTree)
+
+When the `TimeStructure` is a [`TwoLevelTree`](@ref), `strat_periods` returns a
+[`StratTreeNodes`](@ref) type, which, through iteration, provides [`StratNode`](@ref) types.
 
+These are equivalent of a [`StrategicPeriod`](@ref) of a [`TwoLevel`](@ref) time structure.
 """
-    opscenarios(ts::TwoLevelTree{T,OP}) where {T,OP}
+strat_periods(ts::TwoLevelTree) = StratTreeNodes(ts)
 
-Returns a collection of all operational scenarios for a [`TwoLevelTree`](@ref) time structure.
 """
-function opscenarios(ts::TwoLevelTree{T,OP}) where {T,OP}
+    opscenarios(ts::TwoLevelTree)
+
+When the `TimeStructure` is a [`TwoLevelTree`](@ref), `opscenarios` returns an `Array` of
+all [`StratNodeOperationalScenario`](@ref)s or [`StratNodeReprOpscenario`](@ref)s types,
+dependening on whether the [`TwoLevelTree`](@ref) includes [`RepresentativePeriods`](@ref)
+or not.
+
+These are equivalent of a [`StratOperationalScenario`](@ref) of a [`TwoLevel`](@ref) time
+structure.
+"""
+function opscenarios(ts::TwoLevelTree)
     return collect(
         Iterators.flatten(opscenarios(sp) for sp in strategic_periods(ts)),
     )
@@ -289,11 +305,14 @@ function opscenarios(ts::TwoLevelTree{T,OP}) where {T,OP}
 end
 
 """
-    repr_periods(ts::TwoLevelTree{T,OP})
+    repr_periods(ts::TwoLevelTree)
+
+When the `TimeStructure` is a [`TwoLevelTree`](@ref), `repr_periods` returns an `Array` of
+all [`StratNodeReprPeriod`](@ref)s.
 
-Returns a collection of all representative periods for a [`TwoLevelTree`](@ref) time structure.
+These are equivalent of a [`StratReprPeriod`](@ref) of a [`TwoLevel`](@ref) time structure.
 """
-function repr_periods(ts::TwoLevelTree{T,OP}) where {T,OP}
+function repr_periods(ts::TwoLevelTree)
     return collect(
         Iterators.flatten(repr_periods(sp) for sp in strategic_periods(ts)),
     )
diff --git a/test/Project.toml b/test/Project.toml
index 956df01..1ce0b8b 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -5,9 +5,10 @@ Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 LocalRegistry = "89398ba2-070a-4b16-a995-9893c55d93cf"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TestItemRunner = "f8b46487-2199-4994-9208-9a1283c18c0a"
+TestItems = "1c621080-faea-4a02-84b6-bbd5e436b8fe"
 TimeStruct = "f9ed5ce0-9f41-4eaa-96da-f38ab8df101c"
 TimeZones = "f269a46b-ccf7-5d73-abea-4c690281aa53"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [compat]
-Aqua = "0.7"
\ No newline at end of file
+Aqua = "0.7"
diff --git a/test/runtests.jl b/test/runtests.jl
index 156215a..5c4ef06 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,3 +1,4 @@
+using TestItems
 using TestItemRunner
 using TimeStruct
 using Aqua
@@ -878,14 +879,73 @@ end
     @test sum(sspdiv[t] for t in tsc) == 200 / 0.5
 end
 
-@testitem "TwoLevelTree" begin
+# Function for testing all iterators. Can be beneficial for all other tests as well
+@testmodule TwoLevelTreeTest begin
+    using TimeStruct, Test
+    function fun(regtree, n_sp, n_op; n_sc=nothing, n_rp=nothing)
+        sps = strat_periods(regtree)
+        ops1 = collect(regtree)
+        ops2 = [
+            t for sp in sps for t in sp
+        ]
+        ops3 = [
+            t for sp in sps for sc in opscenarios(sp) for t in sc
+        ]
+        ops4 = [
+            t for sp in sps for sc in repr_periods(sp) for t in sc
+        ]
+        ops5 = [
+            t for sp in sps for rp in repr_periods(sp) for sc in opscenarios(rp) for t in sc
+        ]
+        @test length(ops1) == n_op
+        @test length(ops1) == length(ops2)
+        @test length(ops1) == length(ops3)
+        @test length(ops1) == length(ops4)
+        @test length(ops1) == length(ops5)
+        for (i, op) in enumerate(ops1)
+            @test op == ops2[i]
+            @test op == ops3[i]
+            @test op == ops4[i]
+            @test op == ops5[i]
+        end
+
+        if isnothing(n_rp)
+            n_rp = n_sp
+        end
+        if isnothing(n_sc)
+            n_sc = n_rp
+        end
+
+        @test length(sps) == n_sp
+
+        @test sum(length(repr_periods(sp)) for sp in sps) == n_rp
+        @test sum(
+            length(rp) for rp in repr_periods(regtree)
+        ) == n_op
+        @test sum(
+            length(rp) for sp in sps for rp in repr_periods(sp)
+        ) == n_op
+
+        @test sum(length(opscenarios(sp)) for sp in sps) == n_sc
+        @test sum(
+            length(sc) for sc in opscenarios(regtree)
+        ) == n_op
+        @test sum(
+            length(sc) for sp in sps for sc in opscenarios(sp)
+        ) == n_op
+        @test sum(
+            length(sc) for sp in sps for rp in repr_periods(sp) for sc in opscenarios(rp)
+        ) == n_op
+
+        return ops1
+    end
+end
+
+@testitem "TwoLevelTree with SimpleTimes" setup=[TwoLevelTreeTest] begin
     regtree = TimeStruct.regular_tree(5, [3, 2], SimpleTimes(5, 1))
-    ops = [t for n in TimeStruct.strat_nodes(regtree) for t in n]
-    @test length(ops) == 5 * 10
-    ops2 = collect(regtree)
-    @test ops == ops2
-    ops3 = [t for sp in strategic_periods(regtree) for t in sp]
-    @test ops == ops3
+    n_sp = 10
+    n_op = n_sp * 5
+    ops = TwoLevelTreeTest.fun(regtree, n_sp, n_op)
 
     op = ops[31]
     @test TimeStruct._opscen(op) == 1
@@ -931,123 +991,41 @@ end
     @test dsp[ops[4]] == 5
 end
 
-@testitem "TwoLevelTree and opscenarios" begin
+@testitem "TwoLevelTree with OperationalScenarios" setup=[TwoLevelTreeTest] begin
     regtree = TimeStruct.regular_tree(
         5,
         [3, 2],
         OperationalScenarios(3, SimpleTimes(5, 1)),
     )
-
-    ops1 = collect(regtree)
-    ops2 = [
-        t for sp in strat_periods(regtree) for sc in opscenarios(sp) for t in sc
-    ]
-    @test length(ops1) == length(ops2)
-    for (i, op) in enumerate(ops1)
-        @test op == ops2[i]
-    end
-
-    @test sum(length(opscenarios(sp)) for sp in strat_periods(regtree)) == 30
-    @test sum(
-        length(sc) for sp in strat_periods(regtree) for sc in opscenarios(sp)
-    ) == 150
-
-    sregtree = TimeStruct.regular_tree(5, [3, 2], SimpleTimes(5, 1))
-    ops1 = collect(sregtree)
-    ops2 = [
-        t for sp in strat_periods(sregtree) for sc in opscenarios(sp) for
-        t in sc
-    ]
-    @test length(ops1) == length(ops2)
-    @test ops1 == ops2
-
-    @test sum(length(opscenarios(sp)) for sp in strat_periods(sregtree)) == 10
-    @test sum(
-        length(sc) for sp in strat_periods(sregtree) for sc in opscenarios(sp)
-    ) == 50
+    n_sp = 10
+    n_sc = n_sp * 3
+    n_op = n_sc * 5
+    TwoLevelTreeTest.fun(regtree, n_sp, n_op; n_sc)
 end
 
-@testitem "TwoLevelTree and repr_periods" begin
+@testitem "TwoLevelTree with RepresentativePeriods" setup=[TwoLevelTreeTest] begin
     regtree = TimeStruct.regular_tree(
         5,
         [3, 2],
         RepresentativePeriods(2, 1, SimpleTimes(5, 1)),
     )
-
-    ops1 = collect(regtree)
-    ops2 = [
-        t for sp in strat_periods(regtree) for sc in repr_periods(sp) for t in sc
-    ]
-    @test length(ops1) == length(ops2)
-    for (i, op) in enumerate(ops1)
-        @test op == ops2[i]
-    end
-
-    @test sum(length(repr_periods(sp)) for sp in strat_periods(regtree)) == 20
-    @test sum(
-        length(rper) for sp in strat_periods(regtree) for rper in repr_periods(sp)
-    ) == 100
-
-    sregtree = TimeStruct.regular_tree(5, [3, 2], SimpleTimes(5, 1))
-    ops1 = collect(sregtree)
-    ops2 = [
-        t for sp in strat_periods(sregtree) for rper in repr_periods(sp) for
-        t in rper
-    ]
-    @test length(ops1) == length(ops2)
-    @test ops1 == ops2
-
-    @test sum(length(repr_periods(sp)) for sp in strat_periods(sregtree)) == 10
-    @test sum(
-        length(rper) for sp in strat_periods(sregtree) for rper in repr_periods(sp)
-    ) == 50
+    n_sp = 10
+    n_rp = n_sp * 2
+    n_op = n_rp * 5
+    TwoLevelTreeTest.fun(regtree, n_sp, n_op; n_rp)
 end
 
-@testitem "TwoLevelTree with RepresentativePeriods and OperationalScenarios" begin
+@testitem "TwoLevelTree with RepresentativePeriods and OperationalScenarios" setup=[TwoLevelTreeTest] begin
     regtree = TimeStruct.regular_tree(
         5,
         [3, 2],
         RepresentativePeriods(2, 1, OperationalScenarios(3, SimpleTimes(5, 1))),
     )
-
-    ops1 = collect(regtree)
-    ops2 = [
-        t for sp in strat_periods(regtree) for t in sp
-    ]
-    ops3 = [
-        t for sp in strat_periods(regtree) for sc in opscenarios(sp) for t in sc
-    ]
-    ops4 = [
-        t for sp in strat_periods(regtree) for sc in repr_periods(sp) for t in sc
-    ]
-    ops5 = [
-        t for sp in strat_periods(regtree) for rp in repr_periods(sp) for sc in opscenarios(rp) for t in sc
-    ]
-    @test length(ops1) == length(ops2)
-    @test length(ops1) == length(ops3)
-    @test length(ops1) == length(ops4)
-    @test length(ops1) == length(ops5)
-    for (i, op) in enumerate(ops1)
-        @test op == ops2[i]
-        @test op == ops3[i]
-        @test op == ops4[i]
-        @test op == ops5[i]
-    end
-
-    @test sum(length(repr_periods(sp)) for sp in strat_periods(regtree)) == 20
-    @test sum(length(opscenarios(sp)) for sp in strat_periods(regtree)) == 60
-    @test sum(
-        length(rp) for sp in strat_periods(regtree) for rp in repr_periods(sp)
-    ) == 300
-    @test sum(
-        length(rp) for rp in repr_periods(regtree)
-    ) == 300
-    @test sum(
-        length(sc) for sp in strat_periods(regtree) for sc in opscenarios(sp)
-    ) == 300
-    @test sum(
-        length(sc) for sp in strat_periods(regtree) for rp in repr_periods(sp) for sc in opscenarios(rp)
-    ) == 300
+    n_sp = 10
+    n_rp = n_sp * 2
+    n_sc = n_rp * 3
+    n_op = n_sc * 5
+    TwoLevelTreeTest.fun(regtree, n_sp, n_op; n_sc, n_rp)
 end
 
 @testitem "Strategic scenarios with operational scenarios" begin

From 4b75ac6cf0f650941fb5fa853384c6ae6deeff45 Mon Sep 17 00:00:00 2001
From: Julian Straus <julian.straus@sintef.no>
Date: Wed, 28 Aug 2024 12:00:42 +0200
Subject: [PATCH 07/14] Provided convenient constructors for TreePeriod

---
 src/tree_nodes.jl   | 231 +++++++++++++++++++-------------------------
 src/twoleveltree.jl |  19 ++++
 2 files changed, 120 insertions(+), 130 deletions(-)

diff --git a/src/tree_nodes.jl b/src/tree_nodes.jl
index 6cbb151..cc39a9a 100644
--- a/src/tree_nodes.jl
+++ b/src/tree_nodes.jl
@@ -1,7 +1,7 @@
 """
     AbstractTreeNode{T} <: TimeStructure{T}
 
-Abstract base type for all tree nodes within a [`TwoLevelTree`] type
+Abstract base type for all tree nodes within a [`TwoLevelTree`](@ref) type.
 """
 abstract type AbstractTreeNode{T} <: TimeStructure{T} end
 
@@ -46,16 +46,13 @@ isfirst(n::StratNode) = n.sp == 1
 Base.show(io::IO, n::StratNode) = print(io, "sp$(n.sp)-br$(n.branch)")
 Base.length(n::StratNode) = length(n.operational)
 Base.eltype(::Type{StratNode}) = TreePeriod
-function Base.iterate(itr::StratNode, state = nothing)
+function Base.iterate(n::StratNode, state = nothing)
     next =
-        isnothing(state) ? iterate(itr.operational) :
-        iterate(itr.operational, state)
+        isnothing(state) ? iterate(n.operational) :
+        iterate(n.operational, state)
     next === nothing && return nothing
-    per = next[1]
 
-    mult = itr.mult_sp * multiple(per)
-    return TreePeriod(itr.sp, itr.branch, probability_branch(itr), mult, per),
-    next[2]
+    return TreePeriod(n, next[1]), next[2]
 end
 
 """
@@ -79,45 +76,41 @@ struct StratNodeOperationalScenario{T} <: AbstractOperationalScenario{T}
     operational::TimeStructure{T}
 end
 
-_opscen(os::StratNodeOperationalScenario) = os.scen
-_branch(os::StratNodeOperationalScenario) = os.branch
-_strat_per(os::StratNodeOperationalScenario) = os.sp
+_opscen(osc::StratNodeOperationalScenario) = osc.scen
+_branch(osc::StratNodeOperationalScenario) = osc.branch
+_strat_per(osc::StratNodeOperationalScenario) = osc.sp
 
-probability(os::StratNodeOperationalScenario) = os.prob_branch * prob_scen
-probability_branch(os::StratNodeOperationalScenario) = os.prob_branch
-mult_scen(os::StratNodeOperationalScenario) = os.mult_scen
+probability(osc::StratNodeOperationalScenario) = osc.prob_branch * prob_scen
+probability_branch(osc::StratNodeOperationalScenario) = osc.prob_branch
+mult_scen(osc::StratNodeOperationalScenario) = osc.mult_scen
 
 StrategicTreeIndexable(::Type{<:StratNodeOperationalScenario}) = HasStratTreeIndex()
 StrategicIndexable(::Type{<:StratNodeOperationalScenario}) = HasStratIndex()
 
 # Adding methods to existing Julia functions
-function Base.show(io::IO, os::StratNodeOperationalScenario)
-    return print(io, "sp$(os.sp)-br$(os.branch)-sc$(os.scen)")
+function Base.show(io::IO, osc::StratNodeOperationalScenario)
+    return print(io, "sp$(osc.sp)-br$(osc.branch)-sc$(osc.scen)")
 end
-Base.length(snops::StratNodeOperationalScenario) = length(snops.operational)
+Base.length(osc::StratNodeOperationalScenario) = length(osc.operational)
 Base.eltype(_::StratNodeOperationalScenario) = TreePeriod
-function Base.last(os::StratNodeOperationalScenario)
-    per = last(os.operational)
-    mult = os.mult_sp * multiple(per)
-    return TreePeriod(os.sp, os.branch, probability_branch(os), mult, per)
+function Base.last(osc::StratNodeOperationalScenario)
+    per = last(osc.operational)
+    return TreePeriod(osc, per)
 end
-function Base.getindex(os::StratNodeOperationalScenario, index)
-    per = os.operational[index]
-    mult = os.mult_sp * multiple(per)
-    return TreePeriod(os.sp, os.branch, probability_branch(os), mult, per)
+function Base.getindex(osc::StratNodeOperationalScenario, index)
+    per = osc.operational[index]
+    return TreePeriod(osc, per)
 end
-function Base.eachindex(os::StratNodeOperationalScenario)
-    return eachindex(os.operational)
+function Base.eachindex(osc::StratNodeOperationalScenario)
+    return eachindex(osc.operational)
 end
-function Base.iterate(os::StratNodeOperationalScenario, state = nothing)
+function Base.iterate(osc::StratNodeOperationalScenario, state = nothing)
     next =
-        isnothing(state) ? iterate(os.operational) :
-        iterate(os.operational, state)
+        isnothing(state) ? iterate(osc.operational) :
+        iterate(osc.operational, state)
     isnothing(next) && return nothing
 
-    mult = os.mult_sp * multiple(next[1])
-    return TreePeriod(os.sp, os.branch, os.prob_branch, mult, next[1]),
-    next[2]
+    return TreePeriod(osc, next[1]), next[2]
 end
 
 """
@@ -144,26 +137,26 @@ end
 When the `TimeStructure` is a [`StratNode`](@ref), `opscenarios` returns a
 [`StratNodeOpScens`](@ref) used for iterating through the individual operational scenarios
 """
-function opscenarios(sn::StratNode{S,T,OP}) where {S,T,OP<:TimeStructure{T}}
-    return StratNodeOpScens(sn, opscenarios(sn.operational))
+function opscenarios(n::StratNode{S,T,OP}) where {S,T,OP<:TimeStructure{T}}
+    return StratNodeOpScens(n, opscenarios(n.operational))
 end
 
-Base.length(ops::StratNodeOpScens) = length(ops.opscens)
+Base.length(oscs::StratNodeOpScens) = length(oscs.opscens)
 Base.eltype(_::StratNodeOpScens) = StratNodeOperationalScenario
-function Base.iterate(ops::StratNodeOpScens, state = (nothing, 1))
+function Base.iterate(oscs::StratNodeOpScens, state = (nothing, 1))
     next =
-        isnothing(state[1]) ? iterate(ops.opscens) :
-        iterate(ops.opscens, state[1])
+        isnothing(state[1]) ? iterate(oscs.opscens) :
+        iterate(oscs.opscens, state[1])
     isnothing(next) && return nothing
 
     scen = state[2]
     return StratNodeOperationalScenario(
-        ops.sp,
-        ops.branch,
+        oscs.sp,
+        oscs.branch,
         scen,
-        ops.mult_sp,
+        oscs.mult_sp,
         mult_scen(next[1]),
-        ops.prob_branch,
+        oscs.prob_branch,
         probability(next[1]),
         next[1],
     ),
@@ -191,45 +184,41 @@ struct StratNodeReprPeriod{T,OP<:TimeStructure{T}} <: AbstractRepresentativePeri
 end
 
 
-_rper(srp::StratNodeReprPeriod) = srp.rp
-_branch(srp::StratNodeReprPeriod) = srp.branch
-_strat_per(srp::StratNodeReprPeriod) = srp.sp
+_rper(rp::StratNodeReprPeriod) = rp.rp
+_branch(rp::StratNodeReprPeriod) = rp.branch
+_strat_per(rp::StratNodeReprPeriod) = rp.sp
 
-probability(srp::StratNodeReprPeriod) = srp.prob_branch
-probability_branch(srp::StratNodeReprPeriod) = srp.prob_branch
-multiple(srp::StratNodeReprPeriod, t::OperationalPeriod) = t.multiple / srp.mult_sp
+probability(rp::StratNodeReprPeriod) = rp.prob_branch
+probability_branch(rp::StratNodeReprPeriod) = rp.prob_branch
+multiple(rp::StratNodeReprPeriod, t::OperationalPeriod) = t.multiple / rp.mult_sp
 
 StrategicTreeIndexable(::Type{<:StratNodeReprPeriod}) = HasStratTreeIndex()
 StrategicIndexable(::Type{<:StratNodeReprPeriod}) = HasStratIndex()
 
 # Adding methods to existing Julia functions
-function Base.show(io::IO, srp::StratNodeReprPeriod)
-    return print(io, "sp$(srp.sp)-br$(srp.branch)-sc$(srp.rp)")
+function Base.show(io::IO, rp::StratNodeReprPeriod)
+    return print(io, "sp$(rp.sp)-br$(rp.branch)-sc$(rp.rp)")
 end
 Base.length(snrp::StratNodeReprPeriod) = length(snrp.operational)
 Base.eltype(_::StratNodeReprPeriod) = TreePeriod
-function Base.last(srp::StratNodeReprPeriod)
-    per = last(srp.operational)
-    mult = srp.mult_sp * multiple(per)
-    return TreePeriod(srp.sp, srp.branch, probability_branch(srp), mult, per)
+function Base.last(rp::StratNodeReprPeriod)
+    per = last(rp.operational)
+    return TreePeriod(rp, per)
 end
-function Base.getindex(srp::StratNodeReprPeriod, index)
-    per = srp.operational[index]
-    mult = srp.mult_sp * multiple(per)
-    return TreePeriod(srp.sp, srp.branch, probability_branch(srp), mult, per)
+function Base.getindex(rp::StratNodeReprPeriod, index)
+    per = rp.operational[index]
+    return TreePeriod(rp, per)
 end
-function Base.eachindex(srp::StratNodeReprPeriod)
-    return eachindex(srp.operational)
+function Base.eachindex(rp::StratNodeReprPeriod)
+    return eachindex(rp.operational)
 end
-function Base.iterate(srp::StratNodeReprPeriod, state = nothing)
+function Base.iterate(rp::StratNodeReprPeriod, state = nothing)
     next =
-        isnothing(state) ? iterate(srp.operational) :
-        iterate(srp.operational, state)
+        isnothing(state) ? iterate(rp.operational) :
+        iterate(rp.operational, state)
     isnothing(next) && return nothing
 
-    mult = srp.mult_sp * multiple(next[1])
-    return TreePeriod(srp.sp, srp.branch, srp.prob_branch, mult, next[1]),
-    next[2]
+    return TreePeriod(rp, next[1]), next[2]
 end
 
 """
@@ -255,26 +244,26 @@ end
 
 Iterator that iterates over operational scenarios for a specific strategic node in the tree.
 """
-function repr_periods(sn::StratNode{S,T,OP}) where {S,T,OP<:TimeStructure{T}}
-    return StratNodeReprPeriods(sn, repr_periods(sn.operational))
+function repr_periods(n::StratNode{S,T,OP}) where {S,T,OP<:TimeStructure{T}}
+    return StratNodeReprPeriods(n, repr_periods(n.operational))
 end
 
-Base.length(ops::StratNodeReprPeriods) = length(ops.repr)
+Base.length(rps::StratNodeReprPeriods) = length(rps.repr)
 Base.eltype(_::StratNodeReprPeriods) = StratNodeReprPeriod
-function Base.iterate(ops::StratNodeReprPeriods, state = (nothing, 1))
+function Base.iterate(rps::StratNodeReprPeriods, state = (nothing, 1))
     next =
-        isnothing(state[1]) ? iterate(ops.repr) :
-        iterate(ops.repr, state[1])
+        isnothing(state[1]) ? iterate(rps.repr) :
+        iterate(rps.repr, state[1])
     isnothing(next) && return nothing
 
     scen = state[2]
     return StratNodeReprPeriod(
-        ops.sp,
-        ops.branch,
+        rps.sp,
+        rps.branch,
         scen,
-        ops.mult_sp,
+        rps.mult_sp,
         mult_repr(next[1]),
-        ops.prob_branch,
+        rps.prob_branch,
         next[1],
     ),
     (next[2], scen + 1)
@@ -298,13 +287,13 @@ struct StratNodeReprOpscenario{T} <: AbstractOperationalScenario{T}
     operational::TimeStructure{T}
 end
 
-_opscen(os::StratNodeReprOpscenario) = os.opscen
-_rper(os::StratNodeReprOpscenario) = os.rp
-_branch(os::StratNodeReprOpscenario) = os.branch
-_strat_per(os::StratNodeReprOpscenario) = os.sp
+_opscen(osc::StratNodeReprOpscenario) = osc.opscen
+_rper(osc::StratNodeReprOpscenario) = osc.rp
+_branch(osc::StratNodeReprOpscenario) = osc.branch
+_strat_per(osc::StratNodeReprOpscenario) = osc.sp
 
-probability(os::StratNodeReprOpscenario) = os.prob_branch * os.prob_scen
-probability_branch(os::StratNodeReprOpscenario) = os.prob_branch
+probability(osc::StratNodeReprOpscenario) = osc.prob_branch * osc.prob_scen
+probability_branch(osc::StratNodeReprOpscenario) = osc.prob_branch
 
 StrategicTreeIndexable(::Type{<:StratNodeReprOpscenario}) = HasStratTreeIndex()
 StrategicIndexable(::Type{<:StratNodeReprOpscenario}) = HasStratIndex()
@@ -314,51 +303,33 @@ end
 ScenarioIndexable(::Type{<:StratNodeReprOpscenario}) = HasScenarioIndex()
 
 # Adding methods to existing Julia functions
-function Base.show(io::IO, srop::StratNodeReprOpscenario)
-    return print(io, "sp$(srop.sp)-br$(os.branch)-rp$(srop.rp)-sc$(srop.opscen)")
+function Base.show(io::IO, osc::StratNodeReprOpscenario)
+    return print(io, "sp$(osc.sp)-br$(osc.branch)-rp$(osc.rp)-sc$(osc.opscen)")
 end
-Base.length(srop::StratNodeReprOpscenario) = length(srop.operational)
+Base.length(osc::StratNodeReprOpscenario) = length(osc.operational)
 Base.eltype(_::StratNodeReprOpscenario) = TreePeriod
-function Base.last(os::StratNodeReprOpscenario)
-    per = last(os.operational)
-    rper = ReprPeriod(os.rp, per, os.mult_rp * multiple(per))
-    mult - os.mult_sp * os.mult_rp * multiple(per)
-    return TreePeriod(
-        os.sp,
-        os.branch,
-        probability_branch(os),
-        mult,
-        rper,
-    )
+function Base.last(osc::StratNodeReprOpscenario)
+    per = last(osc.operational)
+    return TreePeriod(osc, per)
 end
 
-function Base.getindex(os::StratNodeReprOpscenario, index)
-    mult = stripunit(os.duration * os.op_per_strat / duration(os.operational))
-    per = ReprPeriod(os.rp, os.operational[index], mult)
-    return TreePeriod(os.sp, os.branch, probability_branch(os), mult, per)
+function Base.getindex(osc::StratNodeReprOpscenario, index)
+    per = osc.operational[index]
+    return TreePeriod(osc, per)
 end
 
-function Base.eachindex(os::StratNodeReprOpscenario)
-    return eachindex(os.operational)
+function Base.eachindex(osc::StratNodeReprOpscenario)
+    return eachindex(osc.operational)
 end
 
 # Iterate the time periods of a StratNodeReprOpscenario
-function Base.iterate(os::StratNodeReprOpscenario, state = nothing)
+function Base.iterate(osc::StratNodeReprOpscenario, state = nothing)
     next =
-        isnothing(state) ? iterate(os.operational) :
-        iterate(os.operational, state)
+        isnothing(state) ? iterate(osc.operational) :
+        iterate(osc.operational, state)
     isnothing(next) && return nothing
 
-    period = ReprPeriod(os.rp, next[1], os.mult_rp * multiple(next[1]))
-    mult = os.mult_sp * os.mult_rp * multiple(next[1])
-    return TreePeriod(
-        os.sp,
-        os.branch,
-        probability_branch(os),
-        mult,
-        period,
-    ),
-    next[2]
+    return TreePeriod(osc, next[1]), next[2]
 end
 
 """
@@ -378,8 +349,8 @@ struct StratNodeReprOpscenarios
     opscens::Any
 end
 
-function StratNodeReprOpscenarios(n::StratNodeReprPeriod{T,OP}, opscens) where {T,OP<:TimeStructure{T}}
-    return StratNodeReprOpscenarios(_strat_per(n), _branch(n), _rper(n), n.mult_sp, n.mult_rp, probability_branch(n), opscens)
+function StratNodeReprOpscenarios(rp::StratNodeReprPeriod{T,OP}, opscens) where {T,OP<:TimeStructure{T}}
+    return StratNodeReprOpscenarios(_strat_per(rp), _branch(rp), _rper(rp), rp.mult_sp, rp.mult_rp, probability_branch(rp), opscens)
 end
 
 """
@@ -388,28 +359,28 @@ end
 Iterator that iterates over operational scenarios for a specific strategic node in the tree.
 """
 function opscenarios(
-    srp::StratNodeReprPeriod{T,RepresentativePeriod{T,OP}},
+    rp::StratNodeReprPeriod{T,RepresentativePeriod{T,OP}},
 ) where {T,OP}
-    return StratNodeReprOpscenarios(srp, opscenarios(srp.operational.operational))
+    return StratNodeReprOpscenarios(rp, opscenarios(rp.operational.operational))
 end
 
-Base.length(srop::StratNodeReprOpscenarios) = length(srop.opscens)
+Base.length(oscs::StratNodeReprOpscenarios) = length(oscs.opscens)
 Base.eltype(_::StratNodeReprOpscenarios) = StratNodeReprOpscenario
-function Base.iterate(srop::StratNodeReprOpscenarios, state = (nothing, 1))
+function Base.iterate(oscs::StratNodeReprOpscenarios, state = (nothing, 1))
     next =
-        isnothing(state[1]) ? iterate(srop.opscens) :
-        iterate(srop.opscens, state[1])
+        isnothing(state[1]) ? iterate(oscs.opscens) :
+        iterate(oscs.opscens, state[1])
     isnothing(next) && return nothing
 
     scen = state[2]
     return StratNodeReprOpscenario(
-        srop.sp,
-        srop.branch,
-        srop.rp,
+        oscs.sp,
+        oscs.branch,
+        oscs.rp,
         scen,
-        srop.mult_sp,
-        srop.mult_rp,
-        srop.prob_branch,
+        oscs.mult_sp,
+        oscs.mult_rp,
+        oscs.prob_branch,
         probability(next[1]),
         next[1],
     ),
diff --git a/src/twoleveltree.jl b/src/twoleveltree.jl
index 9dba329..cad5cb0 100644
--- a/src/twoleveltree.jl
+++ b/src/twoleveltree.jl
@@ -110,6 +110,25 @@ function Base.iterate(itr::TwoLevelTree, state)
     (i, next[2])
 end
 
+# Convenient constructors for the individual types
+function TreePeriod(n::StratNode, per::P) where {P<:Union{TimePeriod, TimeStructure}}
+    mult = n.mult_sp * multiple(per)
+    return TreePeriod(n.sp, n.branch, n.prob_branch, mult, per)
+end
+function TreePeriod(osc::StratNodeOperationalScenario, per::P) where {P<:Union{TimePeriod, AbstractOperationalScenario}}
+    mult = osc.mult_sp * multiple(per)
+    return TreePeriod(osc.sp, osc.branch, osc.prob_branch, mult, per)
+end
+function TreePeriod(rp::StratNodeReprPeriod, per::P) where {P<:Union{TimePeriod, AbstractRepresentativePeriod}}
+    mult = rp.mult_sp * multiple(per)
+    return TreePeriod(rp.sp, rp.branch, rp.prob_branch, mult, per)
+end
+function TreePeriod(osc::StratNodeReprOpscenario, per::P) where {P<:Union{TimePeriod, AbstractOperationalScenario}}
+    rper = ReprPeriod(osc.rp, per, osc.mult_rp * multiple(per))
+    mult = osc.mult_sp * osc.mult_rp * multiple(per)
+    return TreePeriod(osc.sp, osc.branch, osc.prob_branch, mult, rper)
+end
+
 """
     struct StrategicScenario
 

From 4ccab4679b7b10016837b05f2101a2d5fb3b66bc Mon Sep 17 00:00:00 2001
From: Julian Straus <julian.straus@sintef.no>
Date: Wed, 28 Aug 2024 14:40:46 +0200
Subject: [PATCH 08/14] Simplified some iterators

---
 src/tree_nodes.jl   | 102 ++++++++++++++++++++++----------------------
 src/twoleveltree.jl |  34 +++++----------
 2 files changed, 61 insertions(+), 75 deletions(-)

diff --git a/src/tree_nodes.jl b/src/tree_nodes.jl
index cc39a9a..8f15b37 100644
--- a/src/tree_nodes.jl
+++ b/src/tree_nodes.jl
@@ -5,6 +5,23 @@ Abstract base type for all tree nodes within a [`TwoLevelTree`](@ref) type.
 """
 abstract type AbstractTreeNode{T} <: TimeStructure{T} end
 
+"""
+    AbstractTreeStructure
+
+Abstract base type for all tree timestructures within a [`TwoLevelTree`](@ref) type.
+"""
+abstract type AbstractTreeStructure end
+
+Base.length(oscs::AbstractTreeStructure) = length(_oper_struct(oscs))
+function Base.iterate(oscs::AbstractTreeStructure, state = (nothing, 1))
+    next =
+        isnothing(state[1]) ? iterate(_oper_struct(oscs)) :
+        iterate(_oper_struct(oscs), state[1])
+    isnothing(next) && return nothing
+
+    return strat_node_period(oscs, next[1], state[2]), (next[2], state[2] + 1)
+end
+
 abstract type StrategicTreeIndexable end
 struct HasStratTreeIndex <: StrategicTreeIndexable end
 struct NoStratTreeIndex <: StrategicTreeIndexable end
@@ -114,12 +131,12 @@ function Base.iterate(osc::StratNodeOperationalScenario, state = nothing)
 end
 
 """
-    struct StratNodeOpScens
+    struct StratNodeOpScens <: AbstractTreeStructure
 
 Type for iterating through the individual operational scenarios of a [`StratNode`](@ref).
 It is automatically created through the function [`opscenarios`](@ref).
 """
-struct StratNodeOpScens
+struct StratNodeOpScens <: AbstractTreeStructure
     sp::Int
     branch::Int
     mult_sp::Float64
@@ -141,28 +158,21 @@ function opscenarios(n::StratNode{S,T,OP}) where {S,T,OP<:TimeStructure{T}}
     return StratNodeOpScens(n, opscenarios(n.operational))
 end
 
-Base.length(oscs::StratNodeOpScens) = length(oscs.opscens)
-Base.eltype(_::StratNodeOpScens) = StratNodeOperationalScenario
-function Base.iterate(oscs::StratNodeOpScens, state = (nothing, 1))
-    next =
-        isnothing(state[1]) ? iterate(oscs.opscens) :
-        iterate(oscs.opscens, state[1])
-    isnothing(next) && return nothing
-
-    scen = state[2]
+_oper_struct(oscs::StratNodeOpScens) = oscs.opscens
+function strat_node_period(oscs::StratNodeOpScens, next, state)
     return StratNodeOperationalScenario(
         oscs.sp,
         oscs.branch,
-        scen,
+        state,
         oscs.mult_sp,
-        mult_scen(next[1]),
+        mult_scen(next),
         oscs.prob_branch,
-        probability(next[1]),
-        next[1],
-    ),
-    (next[2], scen + 1)
+        probability(next),
+        next,
+    )
 end
 
+Base.eltype(_::StratNodeOpScens) = StratNodeOperationalScenario
 
 """
     struct StratNodeReprPeriod{T,OP<:TimeStructure{T}} <: AbstractRepresentativePeriod{T}
@@ -197,7 +207,7 @@ StrategicIndexable(::Type{<:StratNodeReprPeriod}) = HasStratIndex()
 
 # Adding methods to existing Julia functions
 function Base.show(io::IO, rp::StratNodeReprPeriod)
-    return print(io, "sp$(rp.sp)-br$(rp.branch)-sc$(rp.rp)")
+    return print(io, "sp$(rp.sp)-br$(rp.branch)-rp$(rp.rp)")
 end
 Base.length(snrp::StratNodeReprPeriod) = length(snrp.operational)
 Base.eltype(_::StratNodeReprPeriod) = TreePeriod
@@ -222,12 +232,12 @@ function Base.iterate(rp::StratNodeReprPeriod, state = nothing)
 end
 
 """
-    struct StratNodeReprPeriods
+    struct StratNodeReprPeriods <: AbstractTreeStructure
 
 Type for iterating through the individual presentative periods of a [`StratNode`](@ref).
 It is automatically created through the function [`repr_periods`](@ref).
 """
-struct StratNodeReprPeriods
+struct StratNodeReprPeriods <: AbstractTreeStructure
     sp::Int
     branch::Int
     mult_sp::Float64
@@ -247,28 +257,21 @@ Iterator that iterates over operational scenarios for a specific strategic node
 function repr_periods(n::StratNode{S,T,OP}) where {S,T,OP<:TimeStructure{T}}
     return StratNodeReprPeriods(n, repr_periods(n.operational))
 end
-
-Base.length(rps::StratNodeReprPeriods) = length(rps.repr)
-Base.eltype(_::StratNodeReprPeriods) = StratNodeReprPeriod
-function Base.iterate(rps::StratNodeReprPeriods, state = (nothing, 1))
-    next =
-        isnothing(state[1]) ? iterate(rps.repr) :
-        iterate(rps.repr, state[1])
-    isnothing(next) && return nothing
-
-    scen = state[2]
+_oper_struct(rps::StratNodeReprPeriods) = rps.repr
+function strat_node_period(rps::StratNodeReprPeriods, next, state)
     return StratNodeReprPeriod(
         rps.sp,
         rps.branch,
-        scen,
+        state,
         rps.mult_sp,
-        mult_repr(next[1]),
+        mult_repr(next),
         rps.prob_branch,
-        next[1],
-    ),
-    (next[2], scen + 1)
+        next,
+    )
 end
 
+Base.eltype(_::StratNodeReprPeriods) = StratNodeReprPeriod
+
 """
     struct StratNodeReprOpscenario{T} <: AbstractOperationalScenario{T}
 
@@ -333,13 +336,13 @@ function Base.iterate(osc::StratNodeReprOpscenario, state = nothing)
 end
 
 """
-    struct StratNodeReprOpscenarios
+    struct StratNodeReprOpscenarios <: AbstractTreeStructure
 
 Type for iterating through the individual operational scenarios of a
 [`StratNodeReprPeriod`](@ref). It is automatically created through the function
 [`opscenarios`](@ref).
 """
-struct StratNodeReprOpscenarios
+struct StratNodeReprOpscenarios <: AbstractTreeStructure
     sp::Int
     branch::Int
     rp::Int
@@ -356,7 +359,8 @@ end
 """
     opscenarios(sp::StratNodeReprPeriod{T,RepresentativePeriod{T,OP}})
 
-Iterator that iterates over operational scenarios for a specific strategic node in the tree.
+Iterator that iterates over operational scenarios for a specific representative period of
+a strategic node in the tree.
 """
 function opscenarios(
     rp::StratNodeReprPeriod{T,RepresentativePeriod{T,OP}},
@@ -364,25 +368,19 @@ function opscenarios(
     return StratNodeReprOpscenarios(rp, opscenarios(rp.operational.operational))
 end
 
-Base.length(oscs::StratNodeReprOpscenarios) = length(oscs.opscens)
-Base.eltype(_::StratNodeReprOpscenarios) = StratNodeReprOpscenario
-function Base.iterate(oscs::StratNodeReprOpscenarios, state = (nothing, 1))
-    next =
-        isnothing(state[1]) ? iterate(oscs.opscens) :
-        iterate(oscs.opscens, state[1])
-    isnothing(next) && return nothing
-
-    scen = state[2]
+_oper_struct(oscs::StratNodeReprOpscenarios) = oscs.opscens
+function strat_node_period(oscs::StratNodeReprOpscenarios, next, state)
     return StratNodeReprOpscenario(
         oscs.sp,
         oscs.branch,
         oscs.rp,
-        scen,
+        state,
         oscs.mult_sp,
         oscs.mult_rp,
         oscs.prob_branch,
-        probability(next[1]),
-        next[1],
-    ),
-    (next[2], scen + 1)
+        probability(next),
+        next,
+    )
 end
+
+Base.eltype(_::StratNodeReprOpscenarios) = StratNodeReprOpscenario
diff --git a/src/twoleveltree.jl b/src/twoleveltree.jl
index cad5cb0..aa61168 100644
--- a/src/twoleveltree.jl
+++ b/src/twoleveltree.jl
@@ -79,22 +79,12 @@ function Base.isless(t1::TreePeriod, t2::TreePeriod)
     return t1.period < t2.period
 end
 
-# Iterate through all time periods as OperationalPeriods
-function Base.iterate(itr::TwoLevelTree)
-    spn = itr.nodes[1]
-    next = iterate(spn.operational)
-    next === nothing && return nothing
-    per = next[1]
-
-    mult = _multiple_adj(itr, 1) * multiple(per)
-    return TreePeriod(spn.sp, spn.branch, probability_branch(spn), mult, per),
-    (1, next[2])
-end
-
-function Base.iterate(itr::TwoLevelTree, state)
+function Base.iterate(itr::TwoLevelTree, state = (1, nothing))
     i = state[1]
     spn = itr.nodes[i]
-    next = iterate(spn.operational, state[2])
+    next =
+        isnothing(state[2]) ? iterate(spn.operational) :
+        iterate(spn.operational, state[2])
     if next === nothing
         i = i + 1
         if i > length(itr.nodes)
@@ -255,7 +245,7 @@ function regular_tree(
 end
 
 """
-    struct StratTreeNodes{T, OP}
+    struct StratTreeNodes{T, OP} <: AbstractTreeStructure
 
 Type for iterating through the individual strategic nodes of a [`TwoLevelTree`](@ref).
 It is automatically created through the function [`strat_periods`](@ref), and hence,
@@ -264,7 +254,7 @@ It is automatically created through the function [`strat_periods`](@ref), and he
 Iterating through `StratTreeNodes` using the WithPrev iterator changes the behaviour,
 although the meaining remains unchanged.
 """
-struct StratTreeNodes{T, OP}
+struct StratTreeNodes{T, OP} <: AbstractTreeStructure
     ts::TwoLevelTree{T, OP}
 end
 
@@ -276,12 +266,6 @@ function Base.iterate(stps::StratTreeNodes, state = nothing)
     next == length(stps) + 1 && return nothing
     return stps.ts.nodes[next], next
 end
-# function Base.iterate(stps::StratTreeNodes)
-#     return stps.ts.nodes[state], 1
-# end
-# function Base.iterate(stps::StratTreeNodes, state)
-#     return stps.ts.nodes[state + 1], state + 1
-# end
 
 function Base.iterate(w::WithPrev{StratTreeNodes{T,OP}}) where {T, OP}
     n = iterate(w.itr)
@@ -320,7 +304,11 @@ function opscenarios(ts::TwoLevelTree)
     return collect(
         Iterators.flatten(opscenarios(sp) for sp in strategic_periods(ts)),
     )
-    return opscens
+end
+function opscenarios(ts::TwoLevelTree{T,StratNode{S,T,OP}}) where {S,T,OP<:RepresentativePeriods}
+    return collect(
+        Iterators.flatten(opscenarios(rp) for sp in strategic_periods(ts) for rp in repr_periods(sp)),
+    )
 end
 
 """

From f05d70dda8eb48d8684cb9f03c23800863dcc67c Mon Sep 17 00:00:00 2001
From: Julian Straus <julian.straus@sintef.no>
Date: Wed, 28 Aug 2024 15:25:51 +0200
Subject: [PATCH 09/14] Unified iterators

---
 src/tree_nodes.jl | 99 ++++++++++++++---------------------------------
 1 file changed, 30 insertions(+), 69 deletions(-)

diff --git a/src/tree_nodes.jl b/src/tree_nodes.jl
index 8f15b37..2eed34a 100644
--- a/src/tree_nodes.jl
+++ b/src/tree_nodes.jl
@@ -73,7 +73,7 @@ function Base.iterate(n::StratNode, state = nothing)
 end
 
 """
-    struct StratNodeOperationalScenario{T} <: AbstractOperationalScenario{T}
+    struct StratNodeOperationalScenario{T,OP<:TimeStructure{T}}  <: AbstractOperationalScenario{T}
 
 A structure representing a single operational scenario for a strategic node supporting
 iteration over its time periods. It is created through iterating through
@@ -82,7 +82,7 @@ iteration over its time periods. It is created through iterating through
 It is equivalent to a [`StratOperationalScenario`](@ref) of a [`TwoLevel`](@ref) time
 structure when utilizing a [`TwolevelTree`](@ref).
 """
-struct StratNodeOperationalScenario{T} <: AbstractOperationalScenario{T}
+struct StratNodeOperationalScenario{T,OP<:TimeStructure{T}} <: AbstractOperationalScenario{T}
     sp::Int
     branch::Int
     scen::Int
@@ -90,7 +90,7 @@ struct StratNodeOperationalScenario{T} <: AbstractOperationalScenario{T}
     mult_scen::Float64
     prob_branch::Float64
     prob_scen::Float64
-    operational::TimeStructure{T}
+    operational::OP
 end
 
 _opscen(osc::StratNodeOperationalScenario) = osc.scen
@@ -108,27 +108,7 @@ StrategicIndexable(::Type{<:StratNodeOperationalScenario}) = HasStratIndex()
 function Base.show(io::IO, osc::StratNodeOperationalScenario)
     return print(io, "sp$(osc.sp)-br$(osc.branch)-sc$(osc.scen)")
 end
-Base.length(osc::StratNodeOperationalScenario) = length(osc.operational)
 Base.eltype(_::StratNodeOperationalScenario) = TreePeriod
-function Base.last(osc::StratNodeOperationalScenario)
-    per = last(osc.operational)
-    return TreePeriod(osc, per)
-end
-function Base.getindex(osc::StratNodeOperationalScenario, index)
-    per = osc.operational[index]
-    return TreePeriod(osc, per)
-end
-function Base.eachindex(osc::StratNodeOperationalScenario)
-    return eachindex(osc.operational)
-end
-function Base.iterate(osc::StratNodeOperationalScenario, state = nothing)
-    next =
-        isnothing(state) ? iterate(osc.operational) :
-        iterate(osc.operational, state)
-    isnothing(next) && return nothing
-
-    return TreePeriod(osc, next[1]), next[2]
-end
 
 """
     struct StratNodeOpScens <: AbstractTreeStructure
@@ -209,27 +189,7 @@ StrategicIndexable(::Type{<:StratNodeReprPeriod}) = HasStratIndex()
 function Base.show(io::IO, rp::StratNodeReprPeriod)
     return print(io, "sp$(rp.sp)-br$(rp.branch)-rp$(rp.rp)")
 end
-Base.length(snrp::StratNodeReprPeriod) = length(snrp.operational)
 Base.eltype(_::StratNodeReprPeriod) = TreePeriod
-function Base.last(rp::StratNodeReprPeriod)
-    per = last(rp.operational)
-    return TreePeriod(rp, per)
-end
-function Base.getindex(rp::StratNodeReprPeriod, index)
-    per = rp.operational[index]
-    return TreePeriod(rp, per)
-end
-function Base.eachindex(rp::StratNodeReprPeriod)
-    return eachindex(rp.operational)
-end
-function Base.iterate(rp::StratNodeReprPeriod, state = nothing)
-    next =
-        isnothing(state) ? iterate(rp.operational) :
-        iterate(rp.operational, state)
-    isnothing(next) && return nothing
-
-    return TreePeriod(rp, next[1]), next[2]
-end
 
 """
     struct StratNodeReprPeriods <: AbstractTreeStructure
@@ -278,7 +238,7 @@ Base.eltype(_::StratNodeReprPeriods) = StratNodeReprPeriod
 A structure representing a single operational scenario for a representative period in A
 [`TwoLevelTree`](@ref) structure supporting iteration over its time periods.
 """
-struct StratNodeReprOpscenario{T} <: AbstractOperationalScenario{T}
+struct StratNodeReprOpscenario{T,OP<:TimeStructure{T}}  <: AbstractOperationalScenario{T}
     sp::Int
     branch::Int
     rp::Int
@@ -287,7 +247,7 @@ struct StratNodeReprOpscenario{T} <: AbstractOperationalScenario{T}
     mult_rp::Float64
     prob_branch::Float64
     prob_scen::Float64
-    operational::TimeStructure{T}
+    operational::OP
 end
 
 _opscen(osc::StratNodeReprOpscenario) = osc.opscen
@@ -309,31 +269,7 @@ ScenarioIndexable(::Type{<:StratNodeReprOpscenario}) = HasScenarioIndex()
 function Base.show(io::IO, osc::StratNodeReprOpscenario)
     return print(io, "sp$(osc.sp)-br$(osc.branch)-rp$(osc.rp)-sc$(osc.opscen)")
 end
-Base.length(osc::StratNodeReprOpscenario) = length(osc.operational)
 Base.eltype(_::StratNodeReprOpscenario) = TreePeriod
-function Base.last(osc::StratNodeReprOpscenario)
-    per = last(osc.operational)
-    return TreePeriod(osc, per)
-end
-
-function Base.getindex(osc::StratNodeReprOpscenario, index)
-    per = osc.operational[index]
-    return TreePeriod(osc, per)
-end
-
-function Base.eachindex(osc::StratNodeReprOpscenario)
-    return eachindex(osc.operational)
-end
-
-# Iterate the time periods of a StratNodeReprOpscenario
-function Base.iterate(osc::StratNodeReprOpscenario, state = nothing)
-    next =
-        isnothing(state) ? iterate(osc.operational) :
-        iterate(osc.operational, state)
-    isnothing(next) && return nothing
-
-    return TreePeriod(osc, next[1]), next[2]
-end
 
 """
     struct StratNodeReprOpscenarios <: AbstractTreeStructure
@@ -384,3 +320,28 @@ function strat_node_period(oscs::StratNodeReprOpscenarios, next, state)
 end
 
 Base.eltype(_::StratNodeReprOpscenarios) = StratNodeReprOpscenario
+
+# All introduced subtypes require the same procedures for the iteration and indexing.
+# Hence, all introduced types use the same functions.
+TreeStructure = Union{StratNodeOperationalScenario, StratNodeReprPeriod, StratNodeReprOpscenario}
+Base.length(ts::TreeStructure) = length(ts.operational)
+function Base.last(ts::TreeStructure)
+    per = last(ts.operational)
+    return TreePeriod(ts, per)
+end
+
+function Base.getindex(ts::TreeStructure, index)
+    per = ts.operational[index]
+    return TreePeriod(ts, per)
+end
+function Base.eachindex(ts::TreeStructure)
+    return eachindex(ts.operational)
+end
+function Base.iterate(ts::TreeStructure, state = nothing)
+    next =
+        isnothing(state) ? iterate(ts.operational) :
+        iterate(ts.operational, state)
+    isnothing(next) && return nothing
+
+    return TreePeriod(ts, next[1]), next[2]
+end

From ae7fbea4e12e5dbfb3efe1b58041d29abfe82b0d Mon Sep 17 00:00:00 2001
From: Julian Straus <julian.straus@sintef.no>
Date: Mon, 9 Sep 2024 08:22:03 +0200
Subject: [PATCH 10/14] Moved TreeNodes as subtypes of StrategicPeriods Used
 JuliaFormatter

---
 src/profiles.jl     | 26 +++++++++++----
 src/tree_nodes.jl   | 74 +++++++++++++++++++++++++++++++-----------
 src/twoleveltree.jl | 79 +++++++++++++++++++++++++++++++--------------
 test/runtests.jl    | 54 ++++++++++++++-----------------
 4 files changed, 154 insertions(+), 79 deletions(-)

diff --git a/src/profiles.jl b/src/profiles.jl
index 124a0dd..c31e7a4 100644
--- a/src/profiles.jl
+++ b/src/profiles.jl
@@ -165,7 +165,6 @@ function Base.getindex(
     return _value_lookup(RepresentativeIndexable(T), rp, period)
 end
 
-
 """
     StrategicStochasticProfile(vals::Vector{<:Vector{P}}) where {T<:Duration, P<:TimeProfile{T}}
     StrategicStochasticProfile(vals::Vector{<:Vector{<:Number}})
@@ -187,20 +186,33 @@ profile = StrategicStochasticProfile([
 ])
 ```
 """
-struct StrategicStochasticProfile{T<:Duration,P<:TimeProfile{T}} <: TimeProfile{T}
+struct StrategicStochasticProfile{T<:Duration,P<:TimeProfile{T}} <:
+       TimeProfile{T}
     vals::Vector{<:Vector{P}}
 end
 function StrategicStochasticProfile(vals::Vector{<:Vector{<:Duration}})
-    return StrategicStochasticProfile([[FixedProfile(v_2) for v_2 in v_1] for v_1 in vals])
+    return StrategicStochasticProfile([
+        [FixedProfile(v_2) for v_2 in v_1] for v_1 in vals
+    ])
 end
 
-function _value_lookup(::HasStratTreeIndex, ssp::StrategicStochasticProfile, period)
-    sp_prof = ssp.vals[_strat_per(period) > length(ssp.vals) ? end : _strat_per(period)]
-    branch_prof = sp_prof[_branch(period) > length(sp_prof) ? end : _branch(period)]
+function _value_lookup(
+    ::HasStratTreeIndex,
+    ssp::StrategicStochasticProfile,
+    period,
+)
+    sp_prof = ssp.vals[_strat_per(period) > length(ssp.vals) ? end :
+             _strat_per(period)]
+    branch_prof =
+        sp_prof[_branch(period) > length(sp_prof) ? end : _branch(period)]
     return branch_prof[period]
 end
 
-function _value_lookup(::NoStratTreeIndex, ssp::StrategicStochasticProfile, period)
+function _value_lookup(
+    ::NoStratTreeIndex,
+    ssp::StrategicStochasticProfile,
+    period,
+)
     return error(
         "Type $(typeof(period)) can not be used as index for a strategic stochastic profile",
     )
diff --git a/src/tree_nodes.jl b/src/tree_nodes.jl
index 2eed34a..5a667e6 100644
--- a/src/tree_nodes.jl
+++ b/src/tree_nodes.jl
@@ -1,9 +1,9 @@
 """
-    AbstractTreeNode{T} <: TimeStructure{T}
+    AbstractTreeNode{S,T} <: AbstractStrategicPeriod{S,T}
 
 Abstract base type for all tree nodes within a [`TwoLevelTree`](@ref) type.
 """
-abstract type AbstractTreeNode{T} <: TimeStructure{T} end
+abstract type AbstractTreeNode{S,T} <: AbstractStrategicPeriod{S,T} end
 
 """
     AbstractTreeStructure
@@ -29,10 +29,9 @@ struct NoStratTreeIndex <: StrategicTreeIndexable end
 StrategicTreeIndexable(::Type) = NoStratTreeIndex()
 StrategicTreeIndexable(::Type{<:AbstractTreeNode}) = HasStratTreeIndex()
 StrategicTreeIndexable(::Type{<:TimePeriod}) = HasStratTreeIndex()
-StrategicIndexable(::Type{<:AbstractTreeNode}) = HasStratIndex()
 
 """
-    struct StratNode{S, T, OP<:TimeStructure{T}} <: AbstractTreeNode{T}
+    struct StratNode{S, T, OP<:TimeStructure{T}} <: AbstractTreeNode{S,T}
 
 A structure representing a single strategic node of a [`TwolevelTree`](@ref). It is created
 through iterating through [`StratTreeNodes`](@ref).
@@ -40,7 +39,7 @@ through iterating through [`StratTreeNodes`](@ref).
 It is equivalent to a [`StrategicPeriod`](@ref) of a [`TwoLevel`](@ref) time structure when
 utilizing a [`TwolevelTree`](@ref).
 """
-struct StratNode{S, T, OP<:TimeStructure{T}} <: AbstractTreeNode{T}
+struct StratNode{S,T,OP<:TimeStructure{T}} <: AbstractTreeNode{S,T}
     sp::Int
     branch::Int
     duration::S
@@ -82,7 +81,8 @@ iteration over its time periods. It is created through iterating through
 It is equivalent to a [`StratOperationalScenario`](@ref) of a [`TwoLevel`](@ref) time
 structure when utilizing a [`TwolevelTree`](@ref).
 """
-struct StratNodeOperationalScenario{T,OP<:TimeStructure{T}} <: AbstractOperationalScenario{T}
+struct StratNodeOperationalScenario{T,OP<:TimeStructure{T}} <:
+       AbstractOperationalScenario{T}
     sp::Int
     branch::Int
     scen::Int
@@ -101,7 +101,9 @@ probability(osc::StratNodeOperationalScenario) = osc.prob_branch * prob_scen
 probability_branch(osc::StratNodeOperationalScenario) = osc.prob_branch
 mult_scen(osc::StratNodeOperationalScenario) = osc.mult_scen
 
-StrategicTreeIndexable(::Type{<:StratNodeOperationalScenario}) = HasStratTreeIndex()
+function StrategicTreeIndexable(::Type{<:StratNodeOperationalScenario})
+    return HasStratTreeIndex()
+end
 StrategicIndexable(::Type{<:StratNodeOperationalScenario}) = HasStratIndex()
 
 # Adding methods to existing Julia functions
@@ -124,8 +126,17 @@ struct StratNodeOpScens <: AbstractTreeStructure
     opscens::Any
 end
 
-function StratNodeOpScens(n::StratNode{S,T,OP}, opscens) where {S,T,OP<:TimeStructure{T}}
-    return StratNodeOpScens(_strat_per(n), _branch(n), n.mult_sp, probability_branch(n), opscens)
+function StratNodeOpScens(
+    n::StratNode{S,T,OP},
+    opscens,
+) where {S,T,OP<:TimeStructure{T}}
+    return StratNodeOpScens(
+        _strat_per(n),
+        _branch(n),
+        n.mult_sp,
+        probability_branch(n),
+        opscens,
+    )
 end
 
 """
@@ -163,7 +174,8 @@ A structure representing a single representative period of a [`StrategicNode`](@
 It is equivalent to a [`StratReprPeriod`](@ref) of a [`TwoLevel`](@ref) time structure when
 utilizing a [`TwolevelTree`](@ref).
 """
-struct StratNodeReprPeriod{T,OP<:TimeStructure{T}} <: AbstractRepresentativePeriod{T}
+struct StratNodeReprPeriod{T,OP<:TimeStructure{T}} <:
+       AbstractRepresentativePeriod{T}
     sp::Int
     branch::Int
     rp::Int
@@ -173,14 +185,15 @@ struct StratNodeReprPeriod{T,OP<:TimeStructure{T}} <: AbstractRepresentativePeri
     operational::OP
 end
 
-
 _rper(rp::StratNodeReprPeriod) = rp.rp
 _branch(rp::StratNodeReprPeriod) = rp.branch
 _strat_per(rp::StratNodeReprPeriod) = rp.sp
 
 probability(rp::StratNodeReprPeriod) = rp.prob_branch
 probability_branch(rp::StratNodeReprPeriod) = rp.prob_branch
-multiple(rp::StratNodeReprPeriod, t::OperationalPeriod) = t.multiple / rp.mult_sp
+function multiple(rp::StratNodeReprPeriod, t::OperationalPeriod)
+    return t.multiple / rp.mult_sp
+end
 
 StrategicTreeIndexable(::Type{<:StratNodeReprPeriod}) = HasStratTreeIndex()
 StrategicIndexable(::Type{<:StratNodeReprPeriod}) = HasStratIndex()
@@ -205,8 +218,17 @@ struct StratNodeReprPeriods <: AbstractTreeStructure
     repr::Any
 end
 
-function StratNodeReprPeriods(n::StratNode{S,T,OP}, repr) where {S,T,OP<:TimeStructure{T}}
-    return StratNodeReprPeriods(_strat_per(n), _branch(n), n.mult_sp, probability_branch(n), repr)
+function StratNodeReprPeriods(
+    n::StratNode{S,T,OP},
+    repr,
+) where {S,T,OP<:TimeStructure{T}}
+    return StratNodeReprPeriods(
+        _strat_per(n),
+        _branch(n),
+        n.mult_sp,
+        probability_branch(n),
+        repr,
+    )
 end
 
 """
@@ -238,7 +260,8 @@ Base.eltype(_::StratNodeReprPeriods) = StratNodeReprPeriod
 A structure representing a single operational scenario for a representative period in A
 [`TwoLevelTree`](@ref) structure supporting iteration over its time periods.
 """
-struct StratNodeReprOpscenario{T,OP<:TimeStructure{T}}  <: AbstractOperationalScenario{T}
+struct StratNodeReprOpscenario{T,OP<:TimeStructure{T}} <:
+       AbstractOperationalScenario{T}
     sp::Int
     branch::Int
     rp::Int
@@ -288,8 +311,19 @@ struct StratNodeReprOpscenarios <: AbstractTreeStructure
     opscens::Any
 end
 
-function StratNodeReprOpscenarios(rp::StratNodeReprPeriod{T,OP}, opscens) where {T,OP<:TimeStructure{T}}
-    return StratNodeReprOpscenarios(_strat_per(rp), _branch(rp), _rper(rp), rp.mult_sp, rp.mult_rp, probability_branch(rp), opscens)
+function StratNodeReprOpscenarios(
+    rp::StratNodeReprPeriod{T,OP},
+    opscens,
+) where {T,OP<:TimeStructure{T}}
+    return StratNodeReprOpscenarios(
+        _strat_per(rp),
+        _branch(rp),
+        _rper(rp),
+        rp.mult_sp,
+        rp.mult_rp,
+        probability_branch(rp),
+        opscens,
+    )
 end
 
 """
@@ -323,7 +357,11 @@ Base.eltype(_::StratNodeReprOpscenarios) = StratNodeReprOpscenario
 
 # All introduced subtypes require the same procedures for the iteration and indexing.
 # Hence, all introduced types use the same functions.
-TreeStructure = Union{StratNodeOperationalScenario, StratNodeReprPeriod, StratNodeReprOpscenario}
+TreeStructure = Union{
+    StratNodeOperationalScenario,
+    StratNodeReprPeriod,
+    StratNodeReprOpscenario,
+}
 Base.length(ts::TreeStructure) = length(ts.operational)
 function Base.last(ts::TreeStructure)
     per = last(ts.operational)
diff --git a/src/twoleveltree.jl b/src/twoleveltree.jl
index aa61168..0015104 100644
--- a/src/twoleveltree.jl
+++ b/src/twoleveltree.jl
@@ -7,15 +7,18 @@ 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.
 """
-mutable struct TwoLevelTree{T, OP<:AbstractTreeNode{T}} <: TimeStructure{T}
+mutable struct TwoLevelTree{T,OP<:AbstractTreeNode{T}} <: TimeStructure{T}
     len::Int
     root::Any
     nodes::Vector{OP}
     op_per_strat::Float64
 end
 
-function TwoLevelTree{T, OP}(nodes::Vector{OP}, op_per_strat) where {T, OP<:AbstractTreeNode{T}}
-    return TwoLevelTree{T, OP}(0, nothing, nodes, op_per_strat)
+function TwoLevelTree{T,OP}(
+    nodes::Vector{OP},
+    op_per_strat,
+) where {T,OP<:AbstractTreeNode{T}}
+    return TwoLevelTree{T,OP}(0, nothing, nodes, op_per_strat)
 end
 
 function Base.length(itr::TwoLevelTree)
@@ -31,7 +34,9 @@ function _multiple_adj(itr::TwoLevelTree, n)
 end
 strat_nodes(tree::TwoLevelTree) = tree.nodes
 
-children(n::StratNode, ts::TwoLevelTree) = [c for c in ts.nodes if c.parent == n]
+function children(n::StratNode, ts::TwoLevelTree)
+    return [c for c in ts.nodes if c.parent == n]
+end
 nchildren(n::StratNode, ts::TwoLevelTree) = count(c -> c.parent == n, ts.nodes)
 
 branches(tree::TwoLevelTree, sp) = count(n -> n.sp == sp, tree.nodes)
@@ -101,19 +106,31 @@ function Base.iterate(itr::TwoLevelTree, state = (1, nothing))
 end
 
 # Convenient constructors for the individual types
-function TreePeriod(n::StratNode, per::P) where {P<:Union{TimePeriod, TimeStructure}}
+function TreePeriod(
+    n::StratNode,
+    per::P,
+) where {P<:Union{TimePeriod,TimeStructure}}
     mult = n.mult_sp * multiple(per)
     return TreePeriod(n.sp, n.branch, n.prob_branch, mult, per)
 end
-function TreePeriod(osc::StratNodeOperationalScenario, per::P) where {P<:Union{TimePeriod, AbstractOperationalScenario}}
+function TreePeriod(
+    osc::StratNodeOperationalScenario,
+    per::P,
+) where {P<:Union{TimePeriod,AbstractOperationalScenario}}
     mult = osc.mult_sp * multiple(per)
     return TreePeriod(osc.sp, osc.branch, osc.prob_branch, mult, per)
 end
-function TreePeriod(rp::StratNodeReprPeriod, per::P) where {P<:Union{TimePeriod, AbstractRepresentativePeriod}}
+function TreePeriod(
+    rp::StratNodeReprPeriod,
+    per::P,
+) where {P<:Union{TimePeriod,AbstractRepresentativePeriod}}
     mult = rp.mult_sp * multiple(per)
     return TreePeriod(rp.sp, rp.branch, rp.prob_branch, mult, per)
 end
-function TreePeriod(osc::StratNodeReprOpscenario, per::P) where {P<:Union{TimePeriod, AbstractOperationalScenario}}
+function TreePeriod(
+    osc::StratNodeReprOpscenario,
+    per::P,
+) where {P<:Union{TimePeriod,AbstractOperationalScenario}}
     rper = ReprPeriod(osc.rp, per, osc.mult_rp * multiple(per))
     mult = osc.mult_sp * osc.mult_rp * multiple(per)
     return TreePeriod(osc.sp, osc.branch, osc.prob_branch, mult, rper)
@@ -155,7 +172,6 @@ strategic_scenarios(ts::TwoLevelTree) = StrategicScenarios(ts)
 # TODO: Should be replaced with a single wrapper as it is the case for the other scenarios
 strategic_scenarios(two_level::TwoLevel{S,T}) where {S,T} = [two_level]
 
-
 Base.length(scens::StrategicScenarios) = nleaves(scens.ts)
 function Base.iterate(scs::StrategicScenarios, state = 1)
     if state > nleaves(scs.ts)
@@ -188,18 +204,25 @@ Iterative addition of nodes.
 """
 # Add nodes iteratively in a depth first manner
 function add_node(
-    tree::TwoLevelTree{T, StratNode{S, T, OP}},
+    tree::TwoLevelTree{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)
+) 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}(sp, branches(tree, sp) + 1, duration, prob_branch, mult_sp, parent, oper)
+    node = StratNode{S,T,OP}(
+        sp,
+        branches(tree, sp) + 1,
+        duration,
+        prob_branch,
+        mult_sp,
+        parent,
+        oper,
+    )
     push!(tree.nodes, node)
     if isnothing(parent)
         tree.root = node
@@ -235,9 +258,12 @@ function regular_tree(
     duration::S,
     branching::Vector,
     ts::OP;
-    op_per_strat::Real=1.0,
-) where {S, T, OP<:TimeStructure{T}}
-    tree = TwoLevelTree{T, StratNode{S, T, OP}}(Vector{StratNode{S, T, OP}}(), op_per_strat)
+    op_per_strat::Real = 1.0,
+) where {S,T,OP<:TimeStructure{T}}
+    tree = TwoLevelTree{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)
 
@@ -254,26 +280,26 @@ It is automatically created through the function [`strat_periods`](@ref), and he
 Iterating through `StratTreeNodes` using the WithPrev iterator changes the behaviour,
 although the meaining remains unchanged.
 """
-struct StratTreeNodes{T, OP} <: AbstractTreeStructure
-    ts::TwoLevelTree{T, OP}
+struct StratTreeNodes{T,OP} <: AbstractTreeStructure
+    ts::TwoLevelTree{T,OP}
 end
 
 # Adding methods to existing Julia functions
 Base.length(sps::StratTreeNodes) = length(sps.ts.nodes)
-Base.eltype(_::Type{StratTreeNodes{T, OP}}) where{T, OP} = OP
+Base.eltype(_::Type{StratTreeNodes{T,OP}}) where {T,OP} = OP
 function Base.iterate(stps::StratTreeNodes, state = nothing)
     next = isnothing(state) ? 1 : state + 1
     next == length(stps) + 1 && return nothing
     return stps.ts.nodes[next], next
 end
 
-function Base.iterate(w::WithPrev{StratTreeNodes{T,OP}}) where {T, OP}
+function Base.iterate(w::WithPrev{StratTreeNodes{T,OP}}) where {T,OP}
     n = iterate(w.itr)
     n === nothing && return n
     return (nothing, n[1]), (n[1], n[2])
 end
 
-function Base.iterate(w::WithPrev{StratTreeNodes{T,OP}}, state) where {T, OP}
+function Base.iterate(w::WithPrev{StratTreeNodes{T,OP}}, state) where {T,OP}
     n = iterate(w.itr, state[2])
     n === nothing && return n
     return (n[1].parent, n[1]), (n[1], n[2])
@@ -305,9 +331,14 @@ function opscenarios(ts::TwoLevelTree)
         Iterators.flatten(opscenarios(sp) for sp in strategic_periods(ts)),
     )
 end
-function opscenarios(ts::TwoLevelTree{T,StratNode{S,T,OP}}) where {S,T,OP<:RepresentativePeriods}
+function opscenarios(
+    ts::TwoLevelTree{T,StratNode{S,T,OP}},
+) where {S,T,OP<:RepresentativePeriods}
     return collect(
-        Iterators.flatten(opscenarios(rp) for sp in strategic_periods(ts) for rp in repr_periods(sp)),
+        Iterators.flatten(
+            opscenarios(rp) for sp in strategic_periods(ts) for
+            rp in repr_periods(sp)
+        ),
     )
 end
 
diff --git a/test/runtests.jl b/test/runtests.jl
index 5c4ef06..0fde194 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -882,20 +882,15 @@ end
 # Function for testing all iterators. Can be beneficial for all other tests as well
 @testmodule TwoLevelTreeTest begin
     using TimeStruct, Test
-    function fun(regtree, n_sp, n_op; n_sc=nothing, n_rp=nothing)
+    function fun(regtree, n_sp, n_op; n_sc = nothing, n_rp = nothing)
         sps = strat_periods(regtree)
         ops1 = collect(regtree)
-        ops2 = [
-            t for sp in sps for t in sp
-        ]
-        ops3 = [
-            t for sp in sps for sc in opscenarios(sp) for t in sc
-        ]
-        ops4 = [
-            t for sp in sps for sc in repr_periods(sp) for t in sc
-        ]
+        ops2 = [t for sp in sps for t in sp]
+        ops3 = [t for sp in sps for sc in opscenarios(sp) for t in sc]
+        ops4 = [t for sp in sps for sc in repr_periods(sp) for t in sc]
         ops5 = [
-            t for sp in sps for rp in repr_periods(sp) for sc in opscenarios(rp) for t in sc
+            t for sp in sps for rp in repr_periods(sp) for
+            sc in opscenarios(rp) for t in sc
         ]
         @test length(ops1) == n_op
         @test length(ops1) == length(ops2)
@@ -919,29 +914,22 @@ end
         @test length(sps) == n_sp
 
         @test sum(length(repr_periods(sp)) for sp in sps) == n_rp
-        @test sum(
-            length(rp) for rp in repr_periods(regtree)
-        ) == n_op
-        @test sum(
-            length(rp) for sp in sps for rp in repr_periods(sp)
-        ) == n_op
+        @test sum(length(rp) for rp in repr_periods(regtree)) == n_op
+        @test sum(length(rp) for sp in sps for rp in repr_periods(sp)) == n_op
 
         @test sum(length(opscenarios(sp)) for sp in sps) == n_sc
+        @test sum(length(sc) for sc in opscenarios(regtree)) == n_op
+        @test sum(length(sc) for sp in sps for sc in opscenarios(sp)) == n_op
         @test sum(
-            length(sc) for sc in opscenarios(regtree)
-        ) == n_op
-        @test sum(
-            length(sc) for sp in sps for sc in opscenarios(sp)
-        ) == n_op
-        @test sum(
-            length(sc) for sp in sps for rp in repr_periods(sp) for sc in opscenarios(rp)
+            length(sc) for sp in sps for rp in repr_periods(sp) for
+            sc in opscenarios(rp)
         ) == n_op
 
         return ops1
     end
 end
 
-@testitem "TwoLevelTree with SimpleTimes" setup=[TwoLevelTreeTest] begin
+@testitem "TwoLevelTree with SimpleTimes" setup = [TwoLevelTreeTest] begin
     regtree = TimeStruct.regular_tree(5, [3, 2], SimpleTimes(5, 1))
     n_sp = 10
     n_op = n_sp * 5
@@ -958,7 +946,9 @@ end
 
     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(
+            TimeStruct.probability_branch(n) for n in nodes if n.sp == sp
+        ) ≈ 1.0
     end
     node = nodes[2]
     @test length(node) == 5
@@ -991,7 +981,7 @@ end
     @test dsp[ops[4]] == 5
 end
 
-@testitem "TwoLevelTree with OperationalScenarios" setup=[TwoLevelTreeTest] begin
+@testitem "TwoLevelTree with OperationalScenarios" setup = [TwoLevelTreeTest] begin
     regtree = TimeStruct.regular_tree(
         5,
         [3, 2],
@@ -1003,7 +993,7 @@ end
     TwoLevelTreeTest.fun(regtree, n_sp, n_op; n_sc)
 end
 
-@testitem "TwoLevelTree with RepresentativePeriods" setup=[TwoLevelTreeTest] begin
+@testitem "TwoLevelTree with RepresentativePeriods" setup = [TwoLevelTreeTest] begin
     regtree = TimeStruct.regular_tree(
         5,
         [3, 2],
@@ -1015,7 +1005,8 @@ end
     TwoLevelTreeTest.fun(regtree, n_sp, n_op; n_rp)
 end
 
-@testitem "TwoLevelTree with RepresentativePeriods and OperationalScenarios" setup=[TwoLevelTreeTest] begin
+@testitem "TwoLevelTree with RepresentativePeriods and OperationalScenarios" setup =
+    [TwoLevelTreeTest] begin
     regtree = TimeStruct.regular_tree(
         5,
         [3, 2],
@@ -1054,7 +1045,10 @@ end
 
     scens = TimeStruct.strategic_scenarios(two_level)
     @test length(scens) == 1
-    sps = collect(sp for sc in TimeStruct.strategic_scenarios(two_level) for sp in strat_periods(sc))
+    sps = collect(
+        sp for sc in TimeStruct.strategic_scenarios(two_level) for
+        sp in strat_periods(sc)
+    )
     @test length(sps) == 5
 end
 

From 48d8e52b639c3d8a1267b7cda4e34474476a018d Mon Sep 17 00:00:00 2001
From: Julian Straus <julian.straus@sintef.no>
Date: Wed, 11 Sep 2024 13:16:08 +0200
Subject: [PATCH 11/14] Fixed a bug in opscenarios

---
 src/tree_nodes.jl | 18 +++++++++
 test/runtests.jl  | 94 ++++++++++++++++++++++++++++++-----------------
 2 files changed, 79 insertions(+), 33 deletions(-)

diff --git a/src/tree_nodes.jl b/src/tree_nodes.jl
index 5a667e6..44cc3ae 100644
--- a/src/tree_nodes.jl
+++ b/src/tree_nodes.jl
@@ -148,6 +148,9 @@ When the `TimeStructure` is a [`StratNode`](@ref), `opscenarios` returns a
 function opscenarios(n::StratNode{S,T,OP}) where {S,T,OP<:TimeStructure{T}}
     return StratNodeOpScens(n, opscenarios(n.operational))
 end
+function opscenarios(n::StratNode{S,T,OP}) where {S,T,OP<:RepresentativePeriods}
+    return collect(Iterators.flatten(opscenarios(rp) for rp in repr_periods(n)))
+end
 
 _oper_struct(oscs::StratNodeOpScens) = oscs.opscens
 function strat_node_period(oscs::StratNodeOpScens, next, state)
@@ -335,8 +338,23 @@ a strategic node in the tree.
 function opscenarios(
     rp::StratNodeReprPeriod{T,RepresentativePeriod{T,OP}},
 ) where {T,OP}
+    if _strat_per(rp) == 1 && _rper(rp) == 1
+    end
     return StratNodeReprOpscenarios(rp, opscenarios(rp.operational.operational))
 end
+function opscenarios(
+    rp::StratNodeReprPeriod{T,SingleReprPeriod{T,OP}},
+) where {T,OP}
+    if _strat_per(rp) == 1 && _rper(rp) == 1
+    end
+    return StratNodeOpScens(
+        _strat_per(rp),
+        _branch(rp),
+        rp.mult_sp,
+        probability_branch(rp),
+        opscenarios(rp.operational.ts),
+    )
+end
 
 _oper_struct(oscs::StratNodeReprOpscenarios) = oscs.opscens
 function strat_node_period(oscs::StratNodeReprOpscenarios, next, state)
diff --git a/test/runtests.jl b/test/runtests.jl
index 0fde194..ec06e4d 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -882,50 +882,78 @@ end
 # Function for testing all iterators. Can be beneficial for all other tests as well
 @testmodule TwoLevelTreeTest begin
     using TimeStruct, Test
-    function fun(regtree, n_sp, n_op; n_sc = nothing, n_rp = nothing)
-        sps = strat_periods(regtree)
-        ops1 = collect(regtree)
-        ops2 = [t for sp in sps for t in sp]
-        ops3 = [t for sp in sps for sc in opscenarios(sp) for t in sc]
-        ops4 = [t for sp in sps for sc in repr_periods(sp) for t in sc]
-        ops5 = [
+    function fun(ts, n_sp, n_op; n_sc = nothing, n_rp = nothing)
+        if isnothing(n_rp)
+            n_rp = n_sp
+        end
+        if isnothing(n_sc)
+            n_sc = n_rp
+        end
+
+        # Test the invariants for operational periods
+        sps = strat_periods(ts)
+        ops_inv = Array{Any}(nothing, 8)
+        ops_inv[1] = collect(ts)
+        ops_inv[2] = [t for sp in sps for t in sp]
+
+        ops_inv[3] = [t for rp in repr_periods(ts) for t in rp]
+        ops_inv[4] = [t for sp in sps for rp in repr_periods(sp) for t in rp]
+
+        ops_inv[5] = [t for sc in opscenarios(ts) for t in sc]
+        ops_inv[6] = [t for sp in sps for sc in opscenarios(sp) for t in sc]
+        ops_inv[7] =
+            [t for rp in repr_periods(ts) for sc in opscenarios(rp) for t in sc]
+        ops_inv[8] = [
             t for sp in sps for rp in repr_periods(sp) for
             sc in opscenarios(rp) for t in sc
         ]
-        @test length(ops1) == n_op
-        @test length(ops1) == length(ops2)
-        @test length(ops1) == length(ops3)
-        @test length(ops1) == length(ops4)
-        @test length(ops1) == length(ops5)
-        for (i, op) in enumerate(ops1)
-            @test op == ops2[i]
-            @test op == ops3[i]
-            @test op == ops4[i]
-            @test op == ops5[i]
+
+        for ops in ops_inv
+            @test length(ops) == n_op
+        end
+        for (i, op) in enumerate(ops_inv[1])
+            for ops in ops_inv[2:end]
+                @test op == ops[i]
+            end
         end
 
-        if isnothing(n_rp)
-            n_rp = n_sp
+        # Test the invariants for operational scenarios
+        oscs_inv = Array{Any}(nothing, 4)
+        oscs_inv[1] = [osc for osc in opscenarios(ts)]
+        oscs_inv[2] = [osc for sp in sps for osc in opscenarios(sp)]
+
+        oscs_inv[3] =
+            [osc for rp in repr_periods(ts) for osc in opscenarios(rp)]
+        oscs_inv[4] = [
+            osc for sp in sps for rp in repr_periods(sp) for
+            osc in opscenarios(rp)
+        ]
+
+        for oscs in oscs_inv
+            @test length(oscs) == n_sc
         end
-        if isnothing(n_sc)
-            n_sc = n_rp
+        for (i, osc) in enumerate(oscs_inv[1])
+            for oscs in oscs_inv[2:end]
+                @test osc === oscs[i]
+            end
         end
 
-        @test length(sps) == n_sp
+        # Test the invariants for representative periods
+        rps_inv = Array{Any}(nothing, 2)
 
-        @test sum(length(repr_periods(sp)) for sp in sps) == n_rp
-        @test sum(length(rp) for rp in repr_periods(regtree)) == n_op
-        @test sum(length(rp) for sp in sps for rp in repr_periods(sp)) == n_op
+        rps_inv[1] = [rp for rp in repr_periods(ts)]
+        rps_inv[2] = [rp for sp in sps for rp in repr_periods(sp)]
 
-        @test sum(length(opscenarios(sp)) for sp in sps) == n_sc
-        @test sum(length(sc) for sc in opscenarios(regtree)) == n_op
-        @test sum(length(sc) for sp in sps for sc in opscenarios(sp)) == n_op
-        @test sum(
-            length(sc) for sp in sps for rp in repr_periods(sp) for
-            sc in opscenarios(rp)
-        ) == n_op
+        for rps in rps_inv
+            @test length(rps) == n_rp
+        end
+        for (i, rp) in enumerate(rps_inv[1])
+            for rps in rps_inv[2:end]
+                @test rp == rps[i]
+            end
+        end
 
-        return ops1
+        return ops_inv[1]
     end
 end
 

From 50cc128293dc2f34f76885ecab156ff7c85e327d Mon Sep 17 00:00:00 2001
From: Julian Straus <julian.straus@sintef.no>
Date: Wed, 11 Sep 2024 13:16:43 +0200
Subject: [PATCH 12/14] Minor rework

---
 src/profiles.jl   |  6 ++---
 src/tree_nodes.jl | 61 ++++++++++++++++++++++-------------------------
 2 files changed, 31 insertions(+), 36 deletions(-)

diff --git a/src/profiles.jl b/src/profiles.jl
index c31e7a4..0c90917 100644
--- a/src/profiles.jl
+++ b/src/profiles.jl
@@ -166,7 +166,7 @@ function Base.getindex(
 end
 
 """
-    StrategicStochasticProfile(vals::Vector{<:Vector{P}}) where {T<:Duration, P<:TimeProfile{T}}
+    StrategicStochasticProfile(vals::Vector{<:Vector{P}}) where {T<:Number, P<:TimeProfile{T}}
     StrategicStochasticProfile(vals::Vector{<:Vector{<:Number}})
 
 Time profile with a separate time profile for each strategic node in a [`TwoLevelTree`](@ref)
@@ -186,11 +186,11 @@ profile = StrategicStochasticProfile([
 ])
 ```
 """
-struct StrategicStochasticProfile{T<:Duration,P<:TimeProfile{T}} <:
+struct StrategicStochasticProfile{T<:Number,P<:TimeProfile{T}} <:
        TimeProfile{T}
     vals::Vector{<:Vector{P}}
 end
-function StrategicStochasticProfile(vals::Vector{<:Vector{<:Duration}})
+function StrategicStochasticProfile(vals::Vector{<:Vector{<:Number}})
     return StrategicStochasticProfile([
         [FixedProfile(v_2) for v_2 in v_1] for v_1 in vals
     ])
diff --git a/src/tree_nodes.jl b/src/tree_nodes.jl
index 44cc3ae..162b4b6 100644
--- a/src/tree_nodes.jl
+++ b/src/tree_nodes.jl
@@ -12,14 +12,14 @@ Abstract base type for all tree timestructures within a [`TwoLevelTree`](@ref) t
 """
 abstract type AbstractTreeStructure end
 
-Base.length(oscs::AbstractTreeStructure) = length(_oper_struct(oscs))
-function Base.iterate(oscs::AbstractTreeStructure, state = (nothing, 1))
+Base.length(ats::AbstractTreeStructure) = length(_oper_struct(ats))
+function Base.iterate(ats::AbstractTreeStructure, state = (nothing, 1))
     next =
-        isnothing(state[1]) ? iterate(_oper_struct(oscs)) :
-        iterate(_oper_struct(oscs), state[1])
+        isnothing(state[1]) ? iterate(_oper_struct(ats)) :
+        iterate(_oper_struct(ats), state[1])
     isnothing(next) && return nothing
 
-    return strat_node_period(oscs, next[1], state[2]), (next[2], state[2] + 1)
+    return strat_node_period(ats, next[1], state[2]), (next[2], state[2] + 1)
 end
 
 abstract type StrategicTreeIndexable end
@@ -126,19 +126,6 @@ struct StratNodeOpScens <: AbstractTreeStructure
     opscens::Any
 end
 
-function StratNodeOpScens(
-    n::StratNode{S,T,OP},
-    opscens,
-) where {S,T,OP<:TimeStructure{T}}
-    return StratNodeOpScens(
-        _strat_per(n),
-        _branch(n),
-        n.mult_sp,
-        probability_branch(n),
-        opscens,
-    )
-end
-
 """
     opscenarios(sp::StratNode{S,T,OP})
 
@@ -146,7 +133,13 @@ When the `TimeStructure` is a [`StratNode`](@ref), `opscenarios` returns a
 [`StratNodeOpScens`](@ref) used for iterating through the individual operational scenarios
 """
 function opscenarios(n::StratNode{S,T,OP}) where {S,T,OP<:TimeStructure{T}}
-    return StratNodeOpScens(n, opscenarios(n.operational))
+    return StratNodeOpScens(
+        _strat_per(n),
+        _branch(n),
+        n.mult_sp,
+        probability_branch(n),
+        opscenarios(n.operational),
+    )
 end
 function opscenarios(n::StratNode{S,T,OP}) where {S,T,OP<:RepresentativePeriods}
     return collect(Iterators.flatten(opscenarios(rp) for rp in repr_periods(n)))
@@ -221,27 +214,21 @@ struct StratNodeReprPeriods <: AbstractTreeStructure
     repr::Any
 end
 
-function StratNodeReprPeriods(
-    n::StratNode{S,T,OP},
-    repr,
-) where {S,T,OP<:TimeStructure{T}}
+"""
+    repr_periods(sp::StratNode{S,T,OP})
+
+Iterator that iterates over operational scenarios for a specific strategic node in the tree.
+"""
+function repr_periods(n::StratNode{S,T,OP}) where {S,T,OP<:TimeStructure{T}}
     return StratNodeReprPeriods(
         _strat_per(n),
         _branch(n),
         n.mult_sp,
         probability_branch(n),
-        repr,
+        repr_periods(n.operational),
     )
 end
 
-"""
-    repr_periods(sp::StratNode{S,T,OP})
-
-Iterator that iterates over operational scenarios for a specific strategic node in the tree.
-"""
-function repr_periods(n::StratNode{S,T,OP}) where {S,T,OP<:TimeStructure{T}}
-    return StratNodeReprPeriods(n, repr_periods(n.operational))
-end
 _oper_struct(rps::StratNodeReprPeriods) = rps.repr
 function strat_node_period(rps::StratNodeReprPeriods, next, state)
     return StratNodeReprPeriod(
@@ -340,7 +327,15 @@ function opscenarios(
 ) where {T,OP}
     if _strat_per(rp) == 1 && _rper(rp) == 1
     end
-    return StratNodeReprOpscenarios(rp, opscenarios(rp.operational.operational))
+    return StratNodeReprOpscenarios(
+        _strat_per(rp),
+        _branch(rp),
+        _rper(rp),
+        rp.mult_sp,
+        rp.mult_rp,
+        probability_branch(rp),
+        opscenarios(rp.operational.operational),
+    )
 end
 function opscenarios(
     rp::StratNodeReprPeriod{T,SingleReprPeriod{T,OP}},

From a4295f81a138ca550f9c966720b9a163c47b3ff6 Mon Sep 17 00:00:00 2001
From: Julian Straus <julian.straus@sintef.no>
Date: Wed, 11 Sep 2024 14:54:18 +0200
Subject: [PATCH 13/14] First take on documentation

---
 docs/make.jl              |  8 +++++++
 docs/src/reference/api.md | 18 +++++++++++++---
 src/profiles.jl           |  3 +--
 src/strat_periods.jl      |  6 ++++++
 src/tree_nodes.jl         | 20 +++++++++---------
 src/twoleveltree.jl       | 44 +++++++++++++++++++++++----------------
 6 files changed, 66 insertions(+), 33 deletions(-)

diff --git a/docs/make.jl b/docs/make.jl
index 21ca29c..3d61d26 100644
--- a/docs/make.jl
+++ b/docs/make.jl
@@ -10,8 +10,16 @@ pages = [
         "Discounting" => "manual/discount.md",
     ],
     "API reference" => "reference/api.md",
+    "Internal reference" => "reference/internal.md",
 ]
 
+DocMeta.setdocmeta!(
+    TimeStruct,
+    :DocTestSetup,
+    :(using TimeStruct);
+    recursive = true,
+)
+
 Documenter.makedocs(
     sitename = "TimeStruct",
     format = Documenter.HTML(;
diff --git a/docs/src/reference/api.md b/docs/src/reference/api.md
index 262a056..f74e7eb 100644
--- a/docs/src/reference/api.md
+++ b/docs/src/reference/api.md
@@ -1,7 +1,7 @@
+
 ```@meta
 CurrentModule = TimeStruct
 ```
-
 # Available time structures
 
 ```@docs
@@ -28,13 +28,17 @@ OperationalScenarios
 TwoLevel
 ```
 
+```@docs
+TwoLevelTree
+regular_tree
+```
+
 # Properties of time periods
 
 ```@docs
 duration
 ```
 
-
 ```@docs
 isfirst
 ```
@@ -43,7 +47,6 @@ isfirst
 multiple
 ```
 
-
 ```@docs
 probability
 ```
@@ -60,6 +63,11 @@ opscenarios
 
 ```@docs
 strat_periods
+strategic_periods
+```
+
+```@docs
+strategic_scenarios
 ```
 
 ```@docs
@@ -97,6 +105,10 @@ ScenarioProfile
 StrategicProfile
 ```
 
+```@docs
+StrategicStochasticProfile
+```
+
 # Discounting
 
 ```@docs
diff --git a/src/profiles.jl b/src/profiles.jl
index 0c90917..f5ae669 100644
--- a/src/profiles.jl
+++ b/src/profiles.jl
@@ -186,8 +186,7 @@ profile = StrategicStochasticProfile([
 ])
 ```
 """
-struct StrategicStochasticProfile{T<:Number,P<:TimeProfile{T}} <:
-       TimeProfile{T}
+struct StrategicStochasticProfile{T<:Number,P<:TimeProfile{T}} <: TimeProfile{T}
     vals::Vector{<:Vector{P}}
 end
 function StrategicStochasticProfile(vals::Vector{<:Vector{<:Number}})
diff --git a/src/strat_periods.jl b/src/strat_periods.jl
index 056e2ac..8aa6b6a 100644
--- a/src/strat_periods.jl
+++ b/src/strat_periods.jl
@@ -205,4 +205,10 @@ _strat_per(ssp::SingleStrategicPeriod) = 1
 strat_periods(ts::TimeStructure) = SingleStrategicPeriodWrapper(ts)
 
 # Allow strategic_periods() in addition to strat_periods()
+
+"""
+    strategic_periods(ts)
+
+Internal convenience constructor for [`strat_periods`](@ref).
+"""
 strategic_periods(ts) = strat_periods(ts)
diff --git a/src/tree_nodes.jl b/src/tree_nodes.jl
index 162b4b6..e4a07e1 100644
--- a/src/tree_nodes.jl
+++ b/src/tree_nodes.jl
@@ -1,14 +1,14 @@
 """
     AbstractTreeNode{S,T} <: AbstractStrategicPeriod{S,T}
 
-Abstract base type for all tree nodes within a [`TwoLevelTree`](@ref) type.
+Abstract base type for all tree nodes within a [`TwoLevelTree`] type.
 """
 abstract type AbstractTreeNode{S,T} <: AbstractStrategicPeriod{S,T} end
 
 """
     AbstractTreeStructure
 
-Abstract base type for all tree timestructures within a [`TwoLevelTree`](@ref) type.
+Abstract base type for all tree timestructures within a [`TwoLevelTree`] type.
 """
 abstract type AbstractTreeStructure end
 
@@ -33,11 +33,11 @@ StrategicTreeIndexable(::Type{<:TimePeriod}) = HasStratTreeIndex()
 """
     struct StratNode{S, T, OP<:TimeStructure{T}} <: AbstractTreeNode{S,T}
 
-A structure representing a single strategic node of a [`TwolevelTree`](@ref). It is created
+A structure representing a single strategic node of a [`TwolevelTree`]. It is created
 through iterating through [`StratTreeNodes`](@ref).
 
 It is equivalent to a [`StrategicPeriod`](@ref) of a [`TwoLevel`](@ref) time structure when
-utilizing a [`TwolevelTree`](@ref).
+utilizing a [`TwolevelTree`].
 """
 struct StratNode{S,T,OP<:TimeStructure{T}} <: AbstractTreeNode{S,T}
     sp::Int
@@ -79,7 +79,7 @@ iteration over its time periods. It is created through iterating through
 [`StratNodeOpScens`](@ref).
 
 It is equivalent to a [`StratOperationalScenario`](@ref) of a [`TwoLevel`](@ref) time
-structure when utilizing a [`TwolevelTree`](@ref).
+structure when utilizing a [`TwolevelTree`].
 """
 struct StratNodeOperationalScenario{T,OP<:TimeStructure{T}} <:
        AbstractOperationalScenario{T}
@@ -164,11 +164,11 @@ Base.eltype(_::StratNodeOpScens) = StratNodeOperationalScenario
 """
     struct StratNodeReprPeriod{T,OP<:TimeStructure{T}} <: AbstractRepresentativePeriod{T}
 
-A structure representing a single representative period of a [`StrategicNode`](@ref) of a
-[`TwolevelTree`](@ref). It is created through iterating through [`StratNodeReprPeriods`](@ref).
+A structure representing a single representative period of a [`StratNode`](@ref) of a
+[`TwolevelTree`]. It is created through iterating through [`StratNodeReprPeriods`](@ref).
 
-It is equivalent to a [`StratReprPeriod`](@ref) of a [`TwoLevel`](@ref) time structure when
-utilizing a [`TwolevelTree`](@ref).
+It is equivalent to a [`StratReprPeriod`] of a [`TwoLevel`](@ref) time structure when
+utilizing a [`TwolevelTree`].
 """
 struct StratNodeReprPeriod{T,OP<:TimeStructure{T}} <:
        AbstractRepresentativePeriod{T}
@@ -248,7 +248,7 @@ Base.eltype(_::StratNodeReprPeriods) = StratNodeReprPeriod
     struct StratNodeReprOpscenario{T} <: AbstractOperationalScenario{T}
 
 A structure representing a single operational scenario for a representative period in A
-[`TwoLevelTree`](@ref) structure supporting iteration over its time periods.
+[`TwoLevelTree`] structure supporting iteration over its time periods.
 """
 struct StratNodeReprOpscenario{T,OP<:TimeStructure{T}} <:
        AbstractOperationalScenario{T}
diff --git a/src/twoleveltree.jl b/src/twoleveltree.jl
index 0015104..61c8b6a 100644
--- a/src/twoleveltree.jl
+++ b/src/twoleveltree.jl
@@ -162,15 +162,30 @@ end
     struct StrategicScenarios
 
 Type for iteration through the individual strategic scenarios represented as
-[`StrategicScenario`](@ref).
+[`StrategicScenario`].
 """
 struct StrategicScenarios
     ts::TwoLevelTree
 end
+
+"""
+    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
+"""
+strategic_scenarios(two_level::TwoLevel) = [two_level]
+
+"""
+When the `TimeStructure` is a [`TwoLevelTree`](@ref), `strategic_scenarios` returns the
+iterator `StrategicScenarios`.
+"""
 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
-strategic_scenarios(two_level::TwoLevel{S,T}) where {S,T} = [two_level]
 
 Base.length(scens::StrategicScenarios) = nleaves(scens.ts)
 function Base.iterate(scs::StrategicScenarios, state = 1)
@@ -277,7 +292,7 @@ Type for iterating through the individual strategic nodes of a [`TwoLevelTree`](
 It is automatically created through the function [`strat_periods`](@ref), and hence,
 [`strategic_periods`](@ref).
 
-Iterating through `StratTreeNodes` using the WithPrev iterator changes the behaviour,
+Iterating through `StratTreeNodes` using the `WithPrev` iterator changes the behaviour,
 although the meaining remains unchanged.
 """
 struct StratTreeNodes{T,OP} <: AbstractTreeStructure
@@ -305,30 +320,25 @@ function Base.iterate(w::WithPrev{StratTreeNodes{T,OP}}, state) where {T,OP}
     return (n[1].parent, n[1]), (n[1], n[2])
 end
 
-"""
-    strat_periods(ts::TwoLevelTree)
-
-When the `TimeStructure` is a [`TwoLevelTree`](@ref), `strat_periods` returns a
+"""When the `TimeStructure` is a [`TwoLevelTree`](@ref), `strat_periods` returns a
 [`StratTreeNodes`](@ref) type, which, through iteration, provides [`StratNode`](@ref) types.
 
-These are equivalent of a [`StrategicPeriod`](@ref) of a [`TwoLevel`](@ref) time structure.
+These are equivalent to a [`StrategicPeriod`](@ref) of a [`TwoLevel`](@ref) time structure.
 """
 strat_periods(ts::TwoLevelTree) = StratTreeNodes(ts)
 
 """
-    opscenarios(ts::TwoLevelTree)
-
 When the `TimeStructure` is a [`TwoLevelTree`](@ref), `opscenarios` returns an `Array` of
 all [`StratNodeOperationalScenario`](@ref)s or [`StratNodeReprOpscenario`](@ref)s types,
 dependening on whether the [`TwoLevelTree`](@ref) includes [`RepresentativePeriods`](@ref)
 or not.
 
-These are equivalent of a [`StratOperationalScenario`](@ref) of a [`TwoLevel`](@ref) time
+These are equivalent to a [`StratOperationalScenario`](@ref) of a [`TwoLevel`](@ref) time
 structure.
 """
 function opscenarios(ts::TwoLevelTree)
     return collect(
-        Iterators.flatten(opscenarios(sp) for sp in strategic_periods(ts)),
+        Iterators.flatten(opscenarios(sp) for sp in strat_periods(ts)),
     )
 end
 function opscenarios(
@@ -336,22 +346,20 @@ function opscenarios(
 ) where {S,T,OP<:RepresentativePeriods}
     return collect(
         Iterators.flatten(
-            opscenarios(rp) for sp in strategic_periods(ts) for
+            opscenarios(rp) for sp in strat_periods(ts) for
             rp in repr_periods(sp)
         ),
     )
 end
 
 """
-    repr_periods(ts::TwoLevelTree)
-
 When the `TimeStructure` is a [`TwoLevelTree`](@ref), `repr_periods` returns an `Array` of
-all [`StratNodeReprPeriod`](@ref)s.
+all [`StratNodeReprPeriod`]s.
 
-These are equivalent of a [`StratReprPeriod`](@ref) of a [`TwoLevel`](@ref) time structure.
+These are equivalent to a [`StratReprPeriod`] of a [`TwoLevel`](@ref) time structure.
 """
 function repr_periods(ts::TwoLevelTree)
     return collect(
-        Iterators.flatten(repr_periods(sp) for sp in strategic_periods(ts)),
+        Iterators.flatten(repr_periods(sp) for sp in strat_periods(ts)),
     )
 end

From d5c4813157c225bd399e911fb6eee3a46c9225a9 Mon Sep 17 00:00:00 2001
From: Julian Straus <julian.straus@sintef.no>
Date: Wed, 11 Sep 2024 14:58:15 +0200
Subject: [PATCH 14/14] Add new page which I forgot

---
 docs/src/reference/internal.md | 82 ++++++++++++++++++++++++++++++++++
 1 file changed, 82 insertions(+)
 create mode 100644 docs/src/reference/internal.md

diff --git a/docs/src/reference/internal.md b/docs/src/reference/internal.md
new file mode 100644
index 0000000..310adc0
--- /dev/null
+++ b/docs/src/reference/internal.md
@@ -0,0 +1,82 @@
+# Internal types
+
+```@meta
+CurrentModule = TimeStruct
+```
+
+## Strategic period types ([`TwoLevelTree`](@ref))
+
+### Single types
+
+```@docs
+AbstractTreeNode
+StratNode
+StrategicScenario
+```
+
+### Iterator types
+
+```@docs
+AbstractTreeStructure
+StratTreeNodes
+
+```
+
+## Strategic period types ([`TwoLevel`](@ref))
+
+### Single types
+
+```@docs
+AbstractStrategicPeriod
+StrategicPeriod
+```
+
+### Iterator types
+
+```@docs
+```
+
+## Representative period types
+
+### Single types
+
+```@docs
+AbstractRepresentativePeriod
+RepresentativePeriod
+StratNodeReprPeriod
+```
+
+### Iterator types
+
+```@docs
+StratNodeReprPeriods
+```
+
+## Operational scenarios types
+
+### Single types
+
+```@docs
+OperationalScenario
+StratOperationalScenario
+StratNodeOperationalScenario
+StratNodeReprOpscenario
+```
+
+### Iterator types
+
+```@docs
+StrategicScenarios
+StratNodeOpScens
+StratNodeReprOpscenarios
+```
+
+## Operational period types
+
+```@docs
+TimePeriod
+SimplePeriod
+CalendarPeriod
+OperationalPeriod
+TreePeriod
+```