From 815457c61241ac37c9f8d2e27701eb6756aa9777 Mon Sep 17 00:00:00 2001
From: Julian Straus <julian.straus@sintef.no>
Date: Mon, 27 Oct 2025 15:54:56 +0100
Subject: [PATCH 1/9] Modified iteration utilities for StrategicScenarios

---
 src/strat_scenarios/core_types.jl | 58 ++++++++++++++++++++++---------
 test/runtests.jl                  | 19 ++++++----
 2 files changed, 54 insertions(+), 23 deletions(-)

diff --git a/src/strat_scenarios/core_types.jl b/src/strat_scenarios/core_types.jl
index 7a96750..f5bcbe2 100644
--- a/src/strat_scenarios/core_types.jl
+++ b/src/strat_scenarios/core_types.jl
@@ -193,22 +193,26 @@ function TreePeriod(n::StratNode, per::TimePeriod)
     mult = n.mult_sp * multiple(per)
     return TreePeriod(_strat_per(n), _branch(n), per, mult, probability_branch(n))
 end
+
 """
-    struct StrategicScenario
+    struct StrategicScenario{S,T,OP<:AbstractTreeNode{S,T}} <: TimeStructure{T}
 
 Description of an individual strategic scenario. It includes all strategic nodes
 corresponding to a scenario, including the probability. It can be utilized within a
 decomposition algorithm.
 """
-struct StrategicScenario
+struct StrategicScenario{S,T,OP<:AbstractTreeNode{S,T}} <: TimeStructure{T}
+    scen::Int64
     probability::Float64
-    nodes::Vector{<:StratNode}
+    nodes::Vector{OP}
 end
 
-# Iterate through strategic periods of scenario
+Base.show(io::IO, scen::StrategicScenario) = print(io, "scen$(scen.scen)")
+
+# Add basic functions of iterators
 Base.length(scen::StrategicScenario) = length(scen.nodes)
 Base.last(scen::StrategicScenario) = last(scen.nodes)
-
+Base.eltype(_::Type{StrategicScenario{S,T,OP}}) where {S,T,OP} = OP
 function Base.iterate(scs::StrategicScenario, state = nothing)
     next = isnothing(state) ? iterate(scs.nodes) : iterate(scs.nodes, state)
     isnothing(next) && return nothing
@@ -216,13 +220,13 @@ function Base.iterate(scs::StrategicScenario, state = nothing)
 end
 
 """
-    struct StrategicScenarios
+    struct StrategicScenarios{S,T,OP<:AbstractTreeNode{S,T}} <: TimeStructure{T}
 
 Type for iteration through the individual strategic scenarios represented as
 [`StrategicScenario`](@ref).
 """
-struct StrategicScenarios
-    ts::TwoLevelTree
+struct StrategicScenarios{S,T,OP<:AbstractTreeNode{S,T}} <: TimeStructure{T}
+    ts::TwoLevelTree{S,T,OP}
 end
 
 """
@@ -246,21 +250,41 @@ strategic_scenarios(ts::TwoLevelTree) = StrategicScenarios(ts)
 # Allow a TwoLevel structure to be used as a tree with one scenario
 # TODO: Should be replaced with a single wrapper as it is the case for the other scenarios
 
-Base.length(scens::StrategicScenarios) = n_leaves(scens.ts)
-function Base.iterate(scs::StrategicScenarios, state = 1)
-    if state > n_leaves(scs.ts)
-        return nothing
-    end
-
-    node = get_leaf(scs.ts, state)
+# Provide a constructor to simplify the design
+function StrategicScenario(
+    scs::StrategicScenarios{S,T,OP},
+    scen::Int,
+) where {S,T,OP<:TimeStructure{T}}
+    node = get_leaf(scs.ts, scen)
     prob = probability_branch(node)
-    nodes = [node]
+    nodes = OP[node]
     while !isnothing(_parent(node))
         node = _parent(node)
         pushfirst!(nodes, node)
     end
 
-    return StrategicScenario(prob, nodes), state + 1
+    return StrategicScenario(scen, prob, nodes)
+end
+
+# Add basic functions of iterators
+Base.length(scens::StrategicScenarios) = n_leaves(scens.ts)
+function Base.eltype(_::StrategicScenarios{S,T,OP}) where {S,T,OP<:TimeStructure{T}}
+    return StrategicScenario{S,T,OP}
+end
+function Base.iterate(scs::StrategicScenarios, state = nothing)
+    scen = isnothing(state) ? 1 : state + 1
+    scen > n_leaves(scs.ts) && return nothing
+
+    return StrategicScenario(scs, scen), scen
+end
+function Base.getindex(scs::StrategicScenarios, index::Int)
+    return StrategicScenario(scs, index)
+end
+function Base.eachindex(scs::StrategicScenarios)
+    return Base.OneTo(n_leaves(scs.ts))
+end
+function Base.last(scs::StrategicScenarios)
+    return StrategicScenario(scs, length(scs))
 end
 
 """
diff --git a/test/runtests.jl b/test/runtests.jl
index 1d045f2..2a939b3 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1212,11 +1212,20 @@ end
 
 @testitem "Strategic scenarios with operational scenarios" begin
     regtree = TwoLevelTree(5, [3, 2], OperationalScenarios(3, SimpleTimes(5, 1)))
+    op_type = OperationalScenarios{Int64,SimpleTimes{Int64}}
+    stratnode_type = TimeStruct.StratNode{Int64,Int64,op_type}
+    scens = strategic_scenarios(regtree)
 
-    @test length(TimeStruct.strategic_scenarios(regtree)) == 6
+    @test eltype(scens) == TimeStruct.StrategicScenario{Int64,Int64,stratnode_type}
+    @test length(scens) == 6
 
-    for sc in TimeStruct.strategic_scenarios(regtree)
+    # Tests that should work but are broken due to inequality
+    @test last(scens) != scens[6]
+
+    for (k, sc) in enumerate(scens)
         @test length(sc) == length(collect(sc))
+        @test repr(sc) == "scen$(k)"
+        @test eltype(typeof(sc)) == stratnode_type
 
         for (prev_sp, sp) in withprev(sc)
             if !isnothing(prev_sp)
@@ -1229,11 +1238,9 @@ end
 @testitem "TwoLevel as a tree" begin
     two_level = TwoLevel(5, 10, SimpleTimes(10, 1))
 
-    scens = TimeStruct.strategic_scenarios(two_level)
+    scens = 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 strategic_scenarios(two_level) for sp in strat_periods(sc))
     @test length(sps) == 5
 end
 

From dd2b46f6e8bc9370bf023d59237e5364c21d0291 Mon Sep 17 00:00:00 2001
From: Julian Straus <julian.straus@sintef.no>
Date: Tue, 28 Oct 2025 08:13:07 +0100
Subject: [PATCH 2/9] Moved to NTuple to achieve === equality for indexing

---
 src/strat_scenarios/core_types.jl | 17 +++++++++--------
 test/runtests.jl                  |  6 ++----
 2 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/src/strat_scenarios/core_types.jl b/src/strat_scenarios/core_types.jl
index f5bcbe2..19dd77a 100644
--- a/src/strat_scenarios/core_types.jl
+++ b/src/strat_scenarios/core_types.jl
@@ -201,10 +201,10 @@ Description of an individual strategic scenario. It includes all strategic nodes
 corresponding to a scenario, including the probability. It can be utilized within a
 decomposition algorithm.
 """
-struct StrategicScenario{S,T,OP<:AbstractTreeNode{S,T}} <: TimeStructure{T}
+struct StrategicScenario{S,T,N,OP<:AbstractTreeNode{S,T}} <: TimeStructure{T}
     scen::Int64
     probability::Float64
-    nodes::Vector{OP}
+    nodes::NTuple{N,<:OP}
 end
 
 Base.show(io::IO, scen::StrategicScenario) = print(io, "scen$(scen.scen)")
@@ -212,7 +212,7 @@ Base.show(io::IO, scen::StrategicScenario) = print(io, "scen$(scen.scen)")
 # Add basic functions of iterators
 Base.length(scen::StrategicScenario) = length(scen.nodes)
 Base.last(scen::StrategicScenario) = last(scen.nodes)
-Base.eltype(_::Type{StrategicScenario{S,T,OP}}) where {S,T,OP} = OP
+Base.eltype(_::Type{StrategicScenario{S,T,N,OP}}) where {S,T,N,OP} = OP
 function Base.iterate(scs::StrategicScenario, state = nothing)
     next = isnothing(state) ? iterate(scs.nodes) : iterate(scs.nodes, state)
     isnothing(next) && return nothing
@@ -257,19 +257,20 @@ function StrategicScenario(
 ) where {S,T,OP<:TimeStructure{T}}
     node = get_leaf(scs.ts, scen)
     prob = probability_branch(node)
-    nodes = OP[node]
-    while !isnothing(_parent(node))
+    n_strat_per = _strat_per(node)
+    nodes = Vector{OP}(undef, n_strat_per)
+    for sp in n_strat_per:-1:1
+        nodes[sp] = node
         node = _parent(node)
-        pushfirst!(nodes, node)
     end
 
-    return StrategicScenario(scen, prob, nodes)
+    return StrategicScenario(scen, prob, Tuple(nodes))
 end
 
 # Add basic functions of iterators
 Base.length(scens::StrategicScenarios) = n_leaves(scens.ts)
 function Base.eltype(_::StrategicScenarios{S,T,OP}) where {S,T,OP<:TimeStructure{T}}
-    return StrategicScenario{S,T,OP}
+    return StrategicScenario
 end
 function Base.iterate(scs::StrategicScenarios, state = nothing)
     scen = isnothing(state) ? 1 : state + 1
diff --git a/test/runtests.jl b/test/runtests.jl
index 2a939b3..648386f 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1216,11 +1216,9 @@ end
     stratnode_type = TimeStruct.StratNode{Int64,Int64,op_type}
     scens = strategic_scenarios(regtree)
 
-    @test eltype(scens) == TimeStruct.StrategicScenario{Int64,Int64,stratnode_type}
+    @test eltype(scens) == TimeStruct.StrategicScenario
     @test length(scens) == 6
-
-    # Tests that should work but are broken due to inequality
-    @test last(scens) != scens[6]
+    @test last(scens) == scens[6]
 
     for (k, sc) in enumerate(scens)
         @test length(sc) == length(collect(sc))

From e1f65c0e22b1e3a388b63392d06f6633308c673c Mon Sep 17 00:00:00 2001
From: Julian Straus <julian.straus@sintef.no>
Date: Tue, 28 Oct 2025 09:48:51 +0100
Subject: [PATCH 3/9] Moved strategic scenarios to separate folder

---
 docs/src/reference/internal.md         |   4 +-
 src/TimeStruct.jl                      |   1 +
 src/strat_scenarios/core_types.jl      |  94 ---------------------
 src/strat_scenarios/strat_scenarios.jl | 110 +++++++++++++++++++++++++
 4 files changed, 114 insertions(+), 95 deletions(-)
 create mode 100644 src/strat_scenarios/strat_scenarios.jl

diff --git a/docs/src/reference/internal.md b/docs/src/reference/internal.md
index 494bfa8..c201f53 100644
--- a/docs/src/reference/internal.md
+++ b/docs/src/reference/internal.md
@@ -14,6 +14,7 @@ TimeStruct.TimeStructOuterIter
 
 ```@docs
 TimeStruct.AbstractTreeNode
+TimeStruct.AbstractStrategicScenario
 TimeStruct.StratNode
 TimeStruct.StrategicScenario
 ```
@@ -22,8 +23,9 @@ TimeStruct.StrategicScenario
 
 ```@docs
 TimeStruct.AbstractTreeStructure
+TimeStruct.AbstractStratScens
 TimeStruct.StratTreeNodes
-TimeStruct.StrategicScenarios
+TimeStruct.StratScens
 ```
 
 ## [Strategic period types ([`TwoLevel`](@ref))](@id int-types-strat_twolevel)
diff --git a/src/TimeStruct.jl b/src/TimeStruct.jl
index 52a8d58..857a2d3 100644
--- a/src/TimeStruct.jl
+++ b/src/TimeStruct.jl
@@ -14,6 +14,7 @@ include("strategic/core_types.jl")
 include("strategic/strat_periods.jl")
 include("strat_scenarios/tree_periods.jl")
 include("strat_scenarios/core_types.jl")
+include("strat_scenarios/strat_scenarios.jl")
 include("representative/rep_periods.jl")
 include("representative/strat_periods.jl")
 include("representative/tree_periods.jl")
diff --git a/src/strat_scenarios/core_types.jl b/src/strat_scenarios/core_types.jl
index 19dd77a..561d465 100644
--- a/src/strat_scenarios/core_types.jl
+++ b/src/strat_scenarios/core_types.jl
@@ -194,100 +194,6 @@ function TreePeriod(n::StratNode, per::TimePeriod)
     return TreePeriod(_strat_per(n), _branch(n), per, mult, probability_branch(n))
 end
 
-"""
-    struct StrategicScenario{S,T,OP<:AbstractTreeNode{S,T}} <: TimeStructure{T}
-
-Description of an individual strategic scenario. It includes all strategic nodes
-corresponding to a scenario, including the probability. It can be utilized within a
-decomposition algorithm.
-"""
-struct StrategicScenario{S,T,N,OP<:AbstractTreeNode{S,T}} <: TimeStructure{T}
-    scen::Int64
-    probability::Float64
-    nodes::NTuple{N,<:OP}
-end
-
-Base.show(io::IO, scen::StrategicScenario) = print(io, "scen$(scen.scen)")
-
-# Add basic functions of iterators
-Base.length(scen::StrategicScenario) = length(scen.nodes)
-Base.last(scen::StrategicScenario) = last(scen.nodes)
-Base.eltype(_::Type{StrategicScenario{S,T,N,OP}}) where {S,T,N,OP} = OP
-function Base.iterate(scs::StrategicScenario, state = nothing)
-    next = isnothing(state) ? iterate(scs.nodes) : iterate(scs.nodes, state)
-    isnothing(next) && return nothing
-    return next[1], next[2]
-end
-
-"""
-    struct StrategicScenarios{S,T,OP<:AbstractTreeNode{S,T}} <: TimeStructure{T}
-
-Type for iteration through the individual strategic scenarios represented as
-[`StrategicScenario`](@ref).
-"""
-struct StrategicScenarios{S,T,OP<:AbstractTreeNode{S,T}} <: TimeStructure{T}
-    ts::TwoLevelTree{S,T,OP}
-end
-
-"""
-    strategic_scenarios(ts::TwoLevel)
-    strategic_scenarios(ts::TwoLevelTree)
-
-This function returns a type for iterating through the individual strategic scenarios of a
-`TwoLevelTree`. The type of the iterator is dependent on the type of the
-input `TimeStructure`.
-
-When the `TimeStructure` is a [`TwoLevel`](@ref), `strategic_scenarios` returns a Vector with
-the `TwoLevel` as a single entry.
-"""
-strategic_scenarios(ts::TwoLevel) = [ts]
-
-"""
-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
-
-# Provide a constructor to simplify the design
-function StrategicScenario(
-    scs::StrategicScenarios{S,T,OP},
-    scen::Int,
-) where {S,T,OP<:TimeStructure{T}}
-    node = get_leaf(scs.ts, scen)
-    prob = probability_branch(node)
-    n_strat_per = _strat_per(node)
-    nodes = Vector{OP}(undef, n_strat_per)
-    for sp in n_strat_per:-1:1
-        nodes[sp] = node
-        node = _parent(node)
-    end
-
-    return StrategicScenario(scen, prob, Tuple(nodes))
-end
-
-# Add basic functions of iterators
-Base.length(scens::StrategicScenarios) = n_leaves(scens.ts)
-function Base.eltype(_::StrategicScenarios{S,T,OP}) where {S,T,OP<:TimeStructure{T}}
-    return StrategicScenario
-end
-function Base.iterate(scs::StrategicScenarios, state = nothing)
-    scen = isnothing(state) ? 1 : state + 1
-    scen > n_leaves(scs.ts) && return nothing
-
-    return StrategicScenario(scs, scen), scen
-end
-function Base.getindex(scs::StrategicScenarios, index::Int)
-    return StrategicScenario(scs, index)
-end
-function Base.eachindex(scs::StrategicScenarios)
-    return Base.OneTo(n_leaves(scs.ts))
-end
-function Base.last(scs::StrategicScenarios)
-    return StrategicScenario(scs, length(scs))
-end
-
 """
     add_node!(
         nodes::Vector{<:StratNode},
diff --git a/src/strat_scenarios/strat_scenarios.jl b/src/strat_scenarios/strat_scenarios.jl
new file mode 100644
index 0000000..765756b
--- /dev/null
+++ b/src/strat_scenarios/strat_scenarios.jl
@@ -0,0 +1,110 @@
+"""
+    abstract type AbstractStrategicScenario{T} <: TimeStructurePeriod{T}
+
+Abstract type used for time structures that represent a strategic scenario.
+These periods are obtained when iterating through the strategic scenarios of a time
+structure declared by the function [`strategic_scenarios`](@ref).
+"""
+abstract type AbstractStrategicScenario{T} <: TimeStructurePeriod{T} end
+
+"""
+    abstract type AbstractStratScens{S,T} <: TimeStructInnerIter
+
+Abstract type used for time structures that represent a collection of strategic scenarios,
+obtained through calling the function [`strategic_scenarios`](@ref).
+"""
+abstract type AbstractStratScens{T} <: TimeStructInnerIter{T} end
+
+"""
+    struct StrategicScenario{S,T,OP<:AbstractTreeNode{S,T}} <: AbstractStrategicScenario{T}
+
+Description of an individual strategic scenario. It includes all strategic nodes
+corresponding to a scenario, including the probability. It can be utilized within a
+decomposition algorithm.
+"""
+struct StrategicScenario{S,T,N,OP<:AbstractTreeNode{S,T}} <: AbstractStrategicScenario{T}
+    scen::Int64
+    probability::Float64
+    nodes::NTuple{N,<:OP}
+end
+
+Base.show(io::IO, scen::StrategicScenario) = print(io, "scen$(scen.scen)")
+
+# Add basic functions of iterators
+Base.length(scen::StrategicScenario) = length(scen.nodes)
+Base.last(scen::StrategicScenario) = last(scen.nodes)
+Base.eltype(_::Type{StrategicScenario{S,T,N,OP}}) where {S,T,N,OP} = OP
+function Base.iterate(scs::StrategicScenario, state = nothing)
+    next = isnothing(state) ? iterate(scs.nodes) : iterate(scs.nodes, state)
+    isnothing(next) && return nothing
+    return next[1], next[2]
+end
+
+"""
+    struct StratScens{S,T,OP<:AbstractTreeNode{S,T}} <: AbstractStratScens{T}
+
+Type for iteration through the individual strategic scenarios represented as
+[`StrategicScenario`](@ref).
+"""
+struct StratScens{S,T,OP<:AbstractTreeNode{S,T}} <: AbstractStratScens{T}
+    ts::TwoLevelTree{S,T,OP}
+end
+
+"""
+    strategic_scenarios(ts::TwoLevel)
+    strategic_scenarios(ts::TwoLevelTree)
+
+This function returns a type for iterating through the individual strategic scenarios of a
+`TwoLevelTree`. The type of the iterator is dependent on the type of the
+input `TimeStructure`.
+
+When the `TimeStructure` is a [`TwoLevel`](@ref), `strategic_scenarios` returns a Vector with
+the `TwoLevel` as a single entry.
+"""
+strategic_scenarios(ts::TwoLevel) = [ts]
+
+"""
+When the `TimeStructure` is a [`TwoLevelTree`](@ref), `strategic_scenarios` returns the
+iterator `StratScens`.
+"""
+strategic_scenarios(ts::TwoLevelTree) = StratScens(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
+
+# Provide a constructor to simplify the design
+function StrategicScenario(
+    scs::StratScens{S,T,OP},
+    scen::Int,
+) where {S,T,OP<:TimeStructure{T}}
+    node = get_leaf(scs.ts, scen)
+    prob = probability_branch(node)
+    n_strat_per = _strat_per(node)
+    nodes = Vector{OP}(undef, n_strat_per)
+    for sp in n_strat_per:-1:1
+        nodes[sp] = node
+        node = _parent(node)
+    end
+
+    return StrategicScenario(scen, prob, Tuple(nodes))
+end
+
+# Add basic functions of iterators
+Base.length(scens::StratScens) = n_leaves(scens.ts)
+function Base.eltype(_::StratScens{S,T,OP}) where {S,T,OP<:TimeStructure{T}}
+    return StrategicScenario
+end
+function Base.iterate(scs::StratScens, state = nothing)
+    scen = isnothing(state) ? 1 : state + 1
+    scen > n_leaves(scs.ts) && return nothing
+
+    return StrategicScenario(scs, scen), scen
+end
+function Base.getindex(scs::StratScens, index::Int)
+    return StrategicScenario(scs, index)
+end
+function Base.eachindex(scs::StratScens)
+    return Base.OneTo(n_leaves(scs.ts))
+end
+function Base.last(scs::StratScens)
+    return StrategicScenario(scs, length(scs))
+end

From 89860259e73ad8a9270ba1f318217d10d3a3c147 Mon Sep 17 00:00:00 2001
From: Julian Straus <julian.straus@sintef.no>
Date: Tue, 28 Oct 2025 10:18:15 +0100
Subject: [PATCH 4/9] Introduced SingleScenarioWrapper

---
 docs/src/reference/internal.md         |  2 +
 src/strat_scenarios/strat_scenarios.jl | 74 ++++++++++++++++++++------
 test/runtests.jl                       | 13 +++++
 3 files changed, 74 insertions(+), 15 deletions(-)

diff --git a/docs/src/reference/internal.md b/docs/src/reference/internal.md
index c201f53..5972578 100644
--- a/docs/src/reference/internal.md
+++ b/docs/src/reference/internal.md
@@ -17,6 +17,7 @@ TimeStruct.AbstractTreeNode
 TimeStruct.AbstractStrategicScenario
 TimeStruct.StratNode
 TimeStruct.StrategicScenario
+TimeStruct.SingleStrategicScenario
 ```
 
 ### [Iterator types](@id int-types-twoleveltree-iter)
@@ -26,6 +27,7 @@ TimeStruct.AbstractTreeStructure
 TimeStruct.AbstractStratScens
 TimeStruct.StratTreeNodes
 TimeStruct.StratScens
+TimeStruct.SingleStrategicScenarioWrapper
 ```
 
 ## [Strategic period types ([`TwoLevel`](@ref))](@id int-types-strat_twolevel)
diff --git a/src/strat_scenarios/strat_scenarios.jl b/src/strat_scenarios/strat_scenarios.jl
index 765756b..4307ab0 100644
--- a/src/strat_scenarios/strat_scenarios.jl
+++ b/src/strat_scenarios/strat_scenarios.jl
@@ -15,6 +15,65 @@ obtained through calling the function [`strategic_scenarios`](@ref).
 """
 abstract type AbstractStratScens{T} <: TimeStructInnerIter{T} end
 
+"""
+    struct SingleStrategicScenario{T,SC<:TimeStructure{T}} <: AbstractStrategicScenario{T}
+
+A type representing a single strategic scenario supporting iteration over its
+time periods. It is created when iterating through [`SingleStrategicScenarioWrapper`](@ref).
+"""
+struct SingleStrategicScenario{T,SC<:TimeStructure{T}} <: AbstractStrategicScenario{T}
+    ts::SC
+end
+
+# Add basic functions of iterators
+Base.length(sc::SingleStrategicScenario) = length(sc.ts)
+Base.eltype(::Type{SingleStrategicScenario{T,SC}}) where {T,SC} = eltype(SC)
+function Base.iterate(sc::SingleStrategicScenario, state = nothing)
+    next = isnothing(state) ? iterate(sc.ts) : iterate(sc.ts, state)
+    return next
+end
+Base.last(sc::SingleStrategicScenario) = last(sc.ts)
+
+"""
+When the `TimeStructure` is a [`SingleStrategicScenario`](@ref), `strat_periods` returns the
+value of its internal [`TimeStructure`](@ref).
+"""
+strat_periods(sc::SingleStrategicScenario) = strat_periods(sc.ts)
+
+"""
+    struct SingleStrategicScenarioWrapper{T,SC<:TimeStructure{T}} <: AbstractStratScens{T}
+
+Type for iterating through the individual strategic periods of a time structure
+without [`TwoLevelTree`](@ref). It is automatically created through the function
+[`strategic_scenarios`](@ref).
+"""
+struct SingleStrategicScenarioWrapper{T,SC<:TimeStructure{T}} <: AbstractStratScens{T}
+    ts::SC
+end
+
+# Add basic functions of iterators
+Base.length(scs::SingleStrategicScenarioWrapper) = 1
+function Base.iterate(scs::SingleStrategicScenarioWrapper, state = nothing)
+    !isnothing(state) && return nothing
+    return SingleStrategicScenario(scs.ts), 1
+end
+function Base.eltype(::Type{SingleStrategicScenarioWrapper{T,SC}}) where {T,SC}
+    return SingleStrategicScenario{T,SC}
+end
+Base.last(scs::SingleStrategicScenarioWrapper) = SingleStrategicScenario(scs.ts)
+
+"""
+    strategic_scenarios(ts::TimeStructure)
+
+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
+[`SingleStrategicScenarioWrapper`](@ref). This corresponds to the default behavior.
+"""
+strategic_scenarios(ts::TimeStructure) = SingleStrategicScenarioWrapper(ts)
+
 """
     struct StrategicScenario{S,T,OP<:AbstractTreeNode{S,T}} <: AbstractStrategicScenario{T}
 
@@ -50,26 +109,11 @@ struct StratScens{S,T,OP<:AbstractTreeNode{S,T}} <: AbstractStratScens{T}
     ts::TwoLevelTree{S,T,OP}
 end
 
-"""
-    strategic_scenarios(ts::TwoLevel)
-    strategic_scenarios(ts::TwoLevelTree)
-
-This function returns a type for iterating through the individual strategic scenarios of a
-`TwoLevelTree`. The type of the iterator is dependent on the type of the
-input `TimeStructure`.
-
-When the `TimeStructure` is a [`TwoLevel`](@ref), `strategic_scenarios` returns a Vector with
-the `TwoLevel` as a single entry.
-"""
-strategic_scenarios(ts::TwoLevel) = [ts]
-
 """
 When the `TimeStructure` is a [`TwoLevelTree`](@ref), `strategic_scenarios` returns the
 iterator `StratScens`.
 """
 strategic_scenarios(ts::TwoLevelTree) = StratScens(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
 
 # Provide a constructor to simplify the design
 function StrategicScenario(
diff --git a/test/runtests.jl b/test/runtests.jl
index 648386f..0c693eb 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1234,10 +1234,23 @@ end
 end
 
 @testitem "TwoLevel as a tree" begin
+    const TS = TimeStruct
     two_level = TwoLevel(5, 10, SimpleTimes(10, 1))
 
+    # Test that we get the correct types and that their utilities are workign
     scens = strategic_scenarios(two_level)
+    @test isa(scens, TS.SingleStrategicScenarioWrapper{Int64, typeof(two_level)})
     @test length(scens) == 1
+    @test eltype(scens) == TimeStrTSuct.SingleStrategicScenario{Int64, typeof(two_level)}
+    @test last(scens) == first(scens)
+
+    scen = first(scens)
+    @test isa(scen, TS.SingleStrategicScenario{Int64, typeof(two_level)})
+    @test length(scen) == 10*5
+    @test eltype(scen) == TS.OperationalPeriod{TS.SimplePeriod{Int64}}
+    @test last(scen) == last(two_level)
+
+    # Test the iterators
     sps = collect(sp for sc in strategic_scenarios(two_level) for sp in strat_periods(sc))
     @test length(sps) == 5
 end

From 2cddda06d1aecb97f047d0e8c7af07b1daaf7627 Mon Sep 17 00:00:00 2001
From: Julian Straus <julian.straus@sintef.no>
Date: Tue, 28 Oct 2025 10:36:26 +0100
Subject: [PATCH 5/9] Added the potential for strat_periods, repr_periods, and
 opscenarios of StrategicScenario and StratScens

---
 src/op_scenarios/strat_periods.jl      | 14 ++------
 src/op_scenarios/tree_periods.jl       | 18 +++--------
 src/representative/tree_periods.jl     |  7 ++--
 src/strat_scenarios/strat_scenarios.jl | 39 +++++++++++++++++++----
 test/runtests.jl                       | 44 +++++++++++++++++++++++---
 5 files changed, 83 insertions(+), 39 deletions(-)

diff --git a/src/op_scenarios/strat_periods.jl b/src/op_scenarios/strat_periods.jl
index d632cc6..8513839 100644
--- a/src/op_scenarios/strat_periods.jl
+++ b/src/op_scenarios/strat_periods.jl
@@ -275,19 +275,9 @@ correct behavior based on the substructure.
 opscenarios(ts::SingleStrategicPeriod) = opscenarios(ts.ts)
 """
 When the `TimeStructure` is a [`TwoLevel`](@ref), `opscenarios` returns a vector of
-[`StratOpScenario`](@ref)s.
+[`StratOpScenario`](@ref)s or a a vector of [`StratReprOpScenario`](@ref)s depending on
+whether the `TimeStructure` includes [`RepresentativePeriods`](@ref) or not.
 """
 function opscenarios(ts::TwoLevel{S,T,OP}) where {S,T,OP}
     return collect(Iterators.flatten(opscenarios(sp) for sp in strategic_periods(ts)))
 end
-"""
-When the `TimeStructure` is a [`TwoLevel`](@ref) with [`RepresentativePeriods`](@ref),
-`opscenarios` returns a vector of [`StratReprOpScenario`](@ref)s.
-"""
-function opscenarios(ts::TwoLevel{S1,T,RepresentativePeriods{S2,T,OP}}) where {S1,S2,T,OP}
-    return collect(
-        Iterators.flatten(
-            opscenarios(rp) for sp in strategic_periods(ts) for rp in repr_periods(sp)
-        ),
-    )
-end
diff --git a/src/op_scenarios/tree_periods.jl b/src/op_scenarios/tree_periods.jl
index a9e7539..4715b46 100644
--- a/src/op_scenarios/tree_periods.jl
+++ b/src/op_scenarios/tree_periods.jl
@@ -239,23 +239,15 @@ function opscenarios(n::StratNode{S,T,OP}) where {S,T,OP<:RepresentativePeriods}
 end
 
 """
-When the `TimeStructure` is a [`TwoLevelTree`](@ref), `opscenarios` returns an `Array` of
-all [`StratNodeOpScenario`](@ref)s or [`StratNodeReprOpScenario`](@ref)s types,
-dependening on whether the [`TwoLevelTree`](@ref) includes [`RepresentativePeriods`](@ref)
+When the `TimeStructure` is a [`TwoLevelTree`](@ref), [`StratScens`](@ref), or a
+[`StrategicScenario`](@ref), `opscenarios` returns an `Array` of all
+[`StratNodeOpScenario`](@ref)s or [`StratNodeReprOpScenario`](@ref)s types,
+depending on whether the `TimeStructure` includes [`RepresentativePeriods`](@ref)
 or not.
 
 These are equivalent to a [`StratOpScenario`](@ref) and [`StratReprOpScenario`](@ref)
 of a [`TwoLevel`](@ref) time structure.
 """
-function opscenarios(ts::TwoLevelTree)
+function opscenarios(ts::Union{TwoLevelTree,StratScens,StrategicScenario})
     return collect(Iterators.flatten(opscenarios(sp) for sp in strat_periods(ts)))
 end
-function opscenarios(
-    ts::TwoLevelTree{T,StratNode{S,T,OP}},
-) where {S,T,OP<:RepresentativePeriods}
-    return collect(
-        Iterators.flatten(
-            opscenarios(rp) for sp in strat_periods(ts) for rp in repr_periods(sp)
-        ),
-    )
-end
diff --git a/src/representative/tree_periods.jl b/src/representative/tree_periods.jl
index 00ebe19..01b219f 100644
--- a/src/representative/tree_periods.jl
+++ b/src/representative/tree_periods.jl
@@ -93,11 +93,12 @@ end
 Base.eltype(_::StratNodeReprPers) = StratNodeReprPeriod
 
 """
-When the `TimeStructure` is a [`TwoLevelTree`](@ref), `repr_periods` returns an `Array` of
-all [`StratNodeReprPeriod`](@ref)s.
+When the `TimeStructure` is a [`TwoLevelTree`](@ref), [`StratScens`](@ref), or a
+[`StrategicScenario`](@ref), `repr_periods` returns an `Array` of all
+[`StratNodeReprPeriod`](@ref)s.
 
 These are equivalent to a [`StratReprPeriod`](@ref) of a [`TwoLevel`](@ref) time structure.
 """
-function repr_periods(ts::TwoLevelTree)
+function repr_periods(ts::Union{TwoLevelTree,StratScens,StrategicScenario})
     return collect(Iterators.flatten(repr_periods(sp) for sp in strat_periods(ts)))
 end
diff --git a/src/strat_scenarios/strat_scenarios.jl b/src/strat_scenarios/strat_scenarios.jl
index 4307ab0..7b3fb78 100644
--- a/src/strat_scenarios/strat_scenarios.jl
+++ b/src/strat_scenarios/strat_scenarios.jl
@@ -34,12 +34,6 @@ function Base.iterate(sc::SingleStrategicScenario, state = nothing)
 end
 Base.last(sc::SingleStrategicScenario) = last(sc.ts)
 
-"""
-When the `TimeStructure` is a [`SingleStrategicScenario`](@ref), `strat_periods` returns the
-value of its internal [`TimeStructure`](@ref).
-"""
-strat_periods(sc::SingleStrategicScenario) = strat_periods(sc.ts)
-
 """
     struct SingleStrategicScenarioWrapper{T,SC<:TimeStructure{T}} <: AbstractStratScens{T}
 
@@ -62,6 +56,14 @@ function Base.eltype(::Type{SingleStrategicScenarioWrapper{T,SC}}) where {T,SC}
 end
 Base.last(scs::SingleStrategicScenarioWrapper) = SingleStrategicScenario(scs.ts)
 
+"""
+When the `TimeStructure` is a [`SingleStrategicScenario`](@ref) or
+[`SingleStrategicScenarioWrapper`](@ref), `strat_periods` returns the value of its internal
+[`TimeStructure`](@ref).
+"""
+strat_periods(sc::SingleStrategicScenario) = strat_periods(sc.ts)
+strat_periods(sc::SingleStrategicScenarioWrapper) = strat_periods(sc.ts)
+
 """
     strategic_scenarios(ts::TimeStructure)
 
@@ -85,6 +87,7 @@ struct StrategicScenario{S,T,N,OP<:AbstractTreeNode{S,T}} <: AbstractStrategicSc
     scen::Int64
     probability::Float64
     nodes::NTuple{N,<:OP}
+    op_per_strat::Float64
 end
 
 Base.show(io::IO, scen::StrategicScenario) = print(io, "scen$(scen.scen)")
@@ -99,6 +102,16 @@ function Base.iterate(scs::StrategicScenario, state = nothing)
     return next[1], next[2]
 end
 
+"""
+When the `TimeStructure` is a [`StrategicScenario`](@ref), `strat_periods` returns a
+[`StratTreeNodes`](@ref) type, which, through iteration, provides [`StratNode`](@ref) types.
+
+These are equivalent to a [`StrategicPeriod`](@ref) of a [`TwoLevel`](@ref) time structure.
+"""
+strat_periods(ts::StrategicScenario) = StratTreeNodes(
+    TwoLevelTree(length(ts), first(ts), [n for n in ts.nodes], ts.op_per_strat),
+)
+
 """
     struct StratScens{S,T,OP<:AbstractTreeNode{S,T}} <: AbstractStratScens{T}
 
@@ -129,7 +142,7 @@ function StrategicScenario(
         node = _parent(node)
     end
 
-    return StrategicScenario(scen, prob, Tuple(nodes))
+    return StrategicScenario(scen, prob, Tuple(nodes), scs.ts.op_per_strat)
 end
 
 # Add basic functions of iterators
@@ -152,3 +165,15 @@ end
 function Base.last(scs::StratScens)
     return StrategicScenario(scs, length(scs))
 end
+
+"""
+When the `TimeStructure` is a [`StratScens`](@ref), `strat_periods` returns a
+[`StratTreeNodes`](@ref) type, which, through iteration, provides [`StratNode`](@ref) types.
+
+These are equivalent to a [`StrategicPeriod`](@ref) of a [`TwoLevel`](@ref) time structure.
+
+!!! note
+    The corresponding `StratTreeNodes` type is equivalent to the created `StratTreeNodes`
+    when using `strat_periods` directly on the [`TwoLevelTree`](@ref).
+"""
+strat_periods(ts::StratScens) = StratTreeNodes(ts.ts)
diff --git a/test/runtests.jl b/test/runtests.jl
index 0c693eb..0dd9232 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1211,15 +1211,49 @@ end
 end
 
 @testitem "Strategic scenarios with operational scenarios" begin
-    regtree = TwoLevelTree(5, [3, 2], OperationalScenarios(3, SimpleTimes(5, 1)))
+    oper_scens = OperationalScenarios(3, SimpleTimes(5, 1))
+    regtree = TwoLevelTree(5, [3, 2], oper_scens)
     op_type = OperationalScenarios{Int64,SimpleTimes{Int64}}
     stratnode_type = TimeStruct.StratNode{Int64,Int64,op_type}
     scens = strategic_scenarios(regtree)
 
+    # Test general functionality
     @test eltype(scens) == TimeStruct.StrategicScenario
     @test length(scens) == 6
     @test last(scens) == scens[6]
 
+    # Test that the strategic periods are correct
+    sps = strat_periods(regtree)
+    sps_scens = strat_periods(scens)
+    sps_scen_1 = strat_periods(first(scens))
+    @test sps == sps_scens
+    @test all(sps1 === sps2 for (sps1, sps2) in zip(sps_scen_1, collect(sps)[1:3]))
+
+    # Test that the representative periods are correct
+    rps = repr_periods(regtree)
+    rps_scens = repr_periods(scens)
+    rps_scen_1 = repr_periods(first(scens))
+    @test rps == rps_scens
+    @test all(rps1 === rps2 for (rps1, rps2) in zip(rps_scen_1, rps[1:9]))
+
+    # Test that the operational scenarios are correct
+    oscs = opscenarios(regtree)
+    oscs_scens = opscenarios(scens)
+    oscs_scen_1 = opscenarios(first(scens))
+    @test oscs == oscs_scens
+    @test all(oscs1 === oscs2 for (oscs1, oscs2) in zip(oscs_scen_1, oscs[1:3]))
+
+    # Test that the operational scenarios are correct when using representative periods
+    rep_pers = RepresentativePeriods(20, [0.25, 0.25, 0.25, 0.25], oper_scens)
+    regtree_rp = TwoLevelTree(5, [3, 2], rep_pers)
+    scens_rp = strategic_scenarios(regtree_rp)
+    oscs_rp = opscenarios(regtree_rp)
+    oscs_rp_scens = opscenarios(scens_rp)
+    oscs_rp_scen_1 = opscenarios(first(scens_rp))
+    @test oscs_rp == oscs_rp_scens
+    @test all(oscs1 === oscs2 for (oscs1, oscs2) in zip(oscs_rp_scen_1, oscs_rp[1:3]))
+
+    # Test some additional functionality
     for (k, sc) in enumerate(scens)
         @test length(sc) == length(collect(sc))
         @test repr(sc) == "scen$(k)"
@@ -1241,7 +1275,7 @@ end
     scens = strategic_scenarios(two_level)
     @test isa(scens, TS.SingleStrategicScenarioWrapper{Int64, typeof(two_level)})
     @test length(scens) == 1
-    @test eltype(scens) == TimeStrTSuct.SingleStrategicScenario{Int64, typeof(two_level)}
+    @test eltype(scens) == TS.SingleStrategicScenario{Int64, typeof(two_level)}
     @test last(scens) == first(scens)
 
     scen = first(scens)
@@ -1251,8 +1285,10 @@ end
     @test last(scen) == last(two_level)
 
     # Test the iterators
-    sps = collect(sp for sc in strategic_scenarios(two_level) for sp in strat_periods(sc))
-    @test length(sps) == 5
+    @test strat_periods(two_level) === strat_periods(scens)
+    @test strat_periods(two_level) === strat_periods(scen)
+    @test length(strat_periods(scens)) == 5
+    @test length(strat_periods(scen)) == 5
 end
 
 @testitem "Profiles constructors" begin

From cb211fc3a63db253ea5a84acef780734b5cdf770 Mon Sep 17 00:00:00 2001
From: Julian Straus <julian.straus@sintef.no>
Date: Wed, 29 Oct 2025 08:47:45 +0100
Subject: [PATCH 6/9] Updated tests for TwoLevelTree

---
 test/runtests.jl | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/test/runtests.jl b/test/runtests.jl
index 0dd9232..1386d7a 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1221,6 +1221,7 @@ end
     @test eltype(scens) == TimeStruct.StrategicScenario
     @test length(scens) == 6
     @test last(scens) == scens[6]
+    @test eachindex(scens) == Base.OneTo(6)
 
     # Test that the strategic periods are correct
     sps = strat_periods(regtree)
@@ -1270,8 +1271,13 @@ end
 @testitem "TwoLevel as a tree" begin
     const TS = TimeStruct
     two_level = TwoLevel(5, 10, SimpleTimes(10, 1))
+    sps = strat_periods(two_level)
 
-    # Test that we get the correct types and that their utilities are workign
+    # Test the Indexing
+    @test TS.StrategicTreeIndexable(typeof(first(sps))) == TS.NoStratTreeIndex()
+    @test TS.StrategicTreeIndexable(typeof(first(first(sps)))) == TS.HasStratTreeIndex()
+
+    # Test that we get the correct types and that their utilities are working
     scens = strategic_scenarios(two_level)
     @test isa(scens, TS.SingleStrategicScenarioWrapper{Int64, typeof(two_level)})
     @test length(scens) == 1
@@ -1284,6 +1290,9 @@ end
     @test eltype(scen) == TS.OperationalPeriod{TS.SimplePeriod{Int64}}
     @test last(scen) == last(two_level)
 
+    # Test that the iteration utilities are working
+    @test all(t_scen == t for (t_scen, t) ∈ zip(scen, two_level))
+
     # Test the iterators
     @test strat_periods(two_level) === strat_periods(scens)
     @test strat_periods(two_level) === strat_periods(scen)

From 3a3b793ba258a5e345380aa081e0eb53f2d1c1e3 Mon Sep 17 00:00:00 2001
From: Julian Straus <julian.straus@sintef.no>
Date: Wed, 29 Oct 2025 15:59:38 +0100
Subject: [PATCH 7/9] Fixed Formatting

---
 test/runtests.jl | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/test/runtests.jl b/test/runtests.jl
index 1386d7a..62ddd1b 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1279,19 +1279,19 @@ end
 
     # Test that we get the correct types and that their utilities are working
     scens = strategic_scenarios(two_level)
-    @test isa(scens, TS.SingleStrategicScenarioWrapper{Int64, typeof(two_level)})
+    @test isa(scens, TS.SingleStrategicScenarioWrapper{Int64,typeof(two_level)})
     @test length(scens) == 1
-    @test eltype(scens) == TS.SingleStrategicScenario{Int64, typeof(two_level)}
+    @test eltype(scens) == TS.SingleStrategicScenario{Int64,typeof(two_level)}
     @test last(scens) == first(scens)
 
     scen = first(scens)
-    @test isa(scen, TS.SingleStrategicScenario{Int64, typeof(two_level)})
-    @test length(scen) == 10*5
+    @test isa(scen, TS.SingleStrategicScenario{Int64,typeof(two_level)})
+    @test length(scen) == 10 * 5
     @test eltype(scen) == TS.OperationalPeriod{TS.SimplePeriod{Int64}}
     @test last(scen) == last(two_level)
 
     # Test that the iteration utilities are working
-    @test all(t_scen == t for (t_scen, t) ∈ zip(scen, two_level))
+    @test all(t_scen == t for (t_scen, t) in zip(scen, two_level))
 
     # Test the iterators
     @test strat_periods(two_level) === strat_periods(scens)

From a14fa5765c9fcac0a6c882533a6a78c7da1e8282 Mon Sep 17 00:00:00 2001
From: Julian Straus <julian.straus@sintef.no>
Date: Fri, 31 Oct 2025 11:20:23 +0100
Subject: [PATCH 8/9] Unified iteration over StrategicScenario

---
 src/strat_scenarios/strat_scenarios.jl | 21 ++++++++++-----
 src/strat_scenarios/tree_periods.jl    |  1 +
 test/runtests.jl                       | 36 ++++++++++++++++----------
 3 files changed, 37 insertions(+), 21 deletions(-)

diff --git a/src/strat_scenarios/strat_scenarios.jl b/src/strat_scenarios/strat_scenarios.jl
index 7b3fb78..a298424 100644
--- a/src/strat_scenarios/strat_scenarios.jl
+++ b/src/strat_scenarios/strat_scenarios.jl
@@ -93,13 +93,20 @@ end
 Base.show(io::IO, scen::StrategicScenario) = print(io, "scen$(scen.scen)")
 
 # Add basic functions of iterators
-Base.length(scen::StrategicScenario) = length(scen.nodes)
-Base.last(scen::StrategicScenario) = last(scen.nodes)
-Base.eltype(_::Type{StrategicScenario{S,T,N,OP}}) where {S,T,N,OP} = OP
-function Base.iterate(scs::StrategicScenario, state = nothing)
-    next = isnothing(state) ? iterate(scs.nodes) : iterate(scs.nodes, state)
-    isnothing(next) && return nothing
-    return next[1], next[2]
+Base.length(scen::StrategicScenario) = sum(length(sn) for sn in scen.nodes)
+Base.last(scen::StrategicScenario) = last(last(scen.nodes))
+Base.eltype(_::Type{StrategicScenario{S,T,N,OP}}) where {S,T,N,OP} = eltype(OP)
+function Base.iterate(scs::StrategicScenario, state = (nothing, 1))
+    sp = state[2]
+    next = isnothing(state[1]) ? iterate(scs.nodes[sp]) : iterate(scs.nodes[sp], state[1])
+    if isnothing(next)
+        sp = sp + 1
+        if sp > length(scs.nodes)
+            return nothing
+        end
+        next = iterate(scs.nodes[sp])
+    end
+    return next[1], (next[2], sp)
 end
 
 """
diff --git a/src/strat_scenarios/tree_periods.jl b/src/strat_scenarios/tree_periods.jl
index b4ca1d8..a91f149 100644
--- a/src/strat_scenarios/tree_periods.jl
+++ b/src/strat_scenarios/tree_periods.jl
@@ -64,6 +64,7 @@ isfirst(n::StratNode) = n.sp == 1
 # 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.last(n::StratNode) = TreePeriod(n, last(n.operational))
 Base.eltype(::Type{StratNode{S,T,OP}}) where {S,T,OP} = TreePeriod{eltype(OP)}
 function Base.iterate(n::StratNode, state = nothing)
     next = isnothing(state) ? iterate(n.operational) : iterate(n.operational, state)
diff --git a/test/runtests.jl b/test/runtests.jl
index 62ddd1b..9771742 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1159,9 +1159,9 @@ end
 
     # Test strategic scenarios
     scens = collect(strategic_scenarios(regtree))
-    @test length(strategic_scenarios(regtree)) == 6
+    @test length(strategic_scenarios(regtree)) == n_leaves(regtree)
     @test length(scens[2].nodes) == regtree.len
-    @test last(scens[1]) == regtree.nodes[3]
+    @test last(scens[1]) == last(regtree.nodes[3])
     @test scens[3].nodes[1] == regtree.nodes[1]
 
     ssp = StrategicStochasticProfile([[10], [11, 12, 13], [20, 21, 22, 23, 30, 40]])
@@ -1210,15 +1210,14 @@ end
     TwoLevelTreeTest.fun(regtree, n_sp, n_op; n_sc, n_rp)
 end
 
-@testitem "Strategic scenarios with operational scenarios" begin
+@testitem "Strategic scenarios with TwoLevelTree" begin
+    const TS = TimeStruct
     oper_scens = OperationalScenarios(3, SimpleTimes(5, 1))
     regtree = TwoLevelTree(5, [3, 2], oper_scens)
-    op_type = OperationalScenarios{Int64,SimpleTimes{Int64}}
-    stratnode_type = TimeStruct.StratNode{Int64,Int64,op_type}
     scens = strategic_scenarios(regtree)
 
     # Test general functionality
-    @test eltype(scens) == TimeStruct.StrategicScenario
+    @test eltype(scens) == TS.StrategicScenario
     @test length(scens) == 6
     @test last(scens) == scens[6]
     @test eachindex(scens) == Base.OneTo(6)
@@ -1254,21 +1253,29 @@ end
     @test oscs_rp == oscs_rp_scens
     @test all(oscs1 === oscs2 for (oscs1, oscs2) in zip(oscs_rp_scen_1, oscs_rp[1:3]))
 
+    # Test that the tree periods are correct
+    ops_scen_1 = first(scens)
+    @test length(ops_scen_1) == 45
+    @test last(ops_scen_1) == collect(regtree)[45]
+    @test all(tp1 === tp2 for (tp1, tp2) in zip(ops_scen_1, collect(regtree)[1:45]))
+    InnerPeriod = TS.ScenarioPeriod{TS.SimplePeriod{Int64}}
+    @test isa(first(ops_scen_1), TS.TreePeriod{InnerPeriod})
+
+    ops_rp_scen_1 = first(scens_rp)
+    @test length(ops_rp_scen_1) == 180
+    @test last(ops_rp_scen_1) == collect(regtree_rp)[180]
+    @test all(tp1 === tp2 for (tp1, tp2) in zip(ops_rp_scen_1, collect(regtree_rp)[1:180]))
+    @test isa(first(ops_rp_scen_1), TS.TreePeriod{TS.ReprPeriod{InnerPeriod}})
+
     # Test some additional functionality
     for (k, sc) in enumerate(scens)
         @test length(sc) == length(collect(sc))
         @test repr(sc) == "scen$(k)"
-        @test eltype(typeof(sc)) == stratnode_type
-
-        for (prev_sp, sp) in withprev(sc)
-            if !isnothing(prev_sp)
-                @test TimeStruct._strat_per(prev_sp) + 1 == TimeStruct._strat_per(sp)
-            end
-        end
+        @test eltype(typeof(sc)) == eltype(typeof(regtree))
     end
 end
 
-@testitem "TwoLevel as a tree" begin
+@testitem "Strategic scenarios with TwoLevel" begin
     const TS = TimeStruct
     two_level = TwoLevel(5, 10, SimpleTimes(10, 1))
     sps = strat_periods(two_level)
@@ -1289,6 +1296,7 @@ end
     @test length(scen) == 10 * 5
     @test eltype(scen) == TS.OperationalPeriod{TS.SimplePeriod{Int64}}
     @test last(scen) == last(two_level)
+    @test all(op1 === op2 for (op1, op2) in zip(scen, two_level))
 
     # Test that the iteration utilities are working
     @test all(t_scen == t for (t_scen, t) in zip(scen, two_level))

From 87ca7ecd42e0dcc0ea6cdcbea5c4a283c907ec15 Mon Sep 17 00:00:00 2001
From: Truls Flatberg <Truls.Flatberg@sintef.no>
Date: Mon, 3 Nov 2025 08:47:33 +0100
Subject: [PATCH 9/9] Align doc with changes

---
 docs/src/manual/multi.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/src/manual/multi.md b/docs/src/manual/multi.md
index 3592de1..01074b2 100644
--- a/docs/src/manual/multi.md
+++ b/docs/src/manual/multi.md
@@ -133,6 +133,6 @@ In the example above, if we only allow one investment in the planning period, th
 
 ```@repl ts
 for sc in strategic_scenarios(two_level_tree)
-    @constraint(m, sum(invest[sp] for sp in sc) <= 1)
+    @constraint(m, sum(invest[sp] for sp in strat_periods(sc)) <= 1)
 end
 ```