Improvements for StrategicScenario

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

docs/src/reference/internal.md

@@ -14,16 +14,20 @@ TimeStruct.TimeStructOuterIter
 
 ```@docs
 TimeStruct.AbstractTreeNode
+TimeStruct.AbstractStrategicScenario
 TimeStruct.StratNode
 TimeStruct.StrategicScenario
+TimeStruct.SingleStrategicScenario
 ```
 
 ### [Iterator types](@id int-types-twoleveltree-iter)
 
 ```@docs
 TimeStruct.AbstractTreeStructure
+TimeStruct.AbstractStratScens
 TimeStruct.StratTreeNodes
-TimeStruct.StrategicScenarios
+TimeStruct.StratScens
+TimeStruct.SingleStrategicScenarioWrapper
 ```
 
 ## [Strategic period types ([`TwoLevel`](@ref))](@id int-types-strat_twolevel)

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

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

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

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

src/strat_scenarios/core_types.jl

@@ -193,75 +193,6 @@ 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
-
-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
-    probability::Float64
-    nodes::Vector{<:StratNode}
-end
-
-# Iterate through strategic periods of scenario
-Base.length(scen::StrategicScenario) = length(scen.nodes)
-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], next[2]
-end
-
-"""
-    struct StrategicScenarios
-
-Type for iteration through the individual strategic scenarios represented as
-[`StrategicScenario`](@ref).
-"""
-struct StrategicScenarios
-    ts::TwoLevelTree
-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
-
-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)
-    prob = probability_branch(node)
-    nodes = [node]
-    while !isnothing(_parent(node))
-        node = _parent(node)
-        pushfirst!(nodes, node)
-    end
-
-    return StrategicScenario(prob, nodes), state + 1
-end
 
 """
     add_node!(

src/strat_scenarios/strat_scenarios.jl

@@ -0,0 +1,186 @@
+"""
+    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 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)
+
+"""
+    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)
+
+"""
+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)
+
+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}
+
+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}
+    op_per_strat::Float64
+end
+
+Base.show(io::IO, scen::StrategicScenario) = print(io, "scen$(scen.scen)")
+
+# Add basic functions of iterators
+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
+
+"""
+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}
+
+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
+
+"""
+When the `TimeStructure` is a [`TwoLevelTree`](@ref), `strategic_scenarios` returns the
+iterator `StratScens`.
+"""
+strategic_scenarios(ts::TwoLevelTree) = StratScens(ts)
+
+# 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), scs.ts.op_per_strat)
+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
+
+"""
+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)

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)