2323# ---------------
2424# StructureData
2525# ---------------
26- if isdefined (:MPI )
27- type MPIWrapper
28- comm:: MPI.Comm
29- init:: Function
30-
31- function MPIWrapper ()
32- instance = new (MPI. Comm (- 1 ))
33- finalizer (instance, freeMPIWrapper)
34-
35- instance. init = function (ucomm:: MPI.Comm )
36- if isdefined (:MPI ) && MPI. Initialized () && ucomm. val == - 1
37- instance. comm = MPI. COMM_WORLD
38- elseif isdefined (:MPI ) && ! MPI. Initialized ()
39- MPI. Init ()
40- instance. comm = MPI. COMM_WORLD
41- elseif isdefined (:MPI ) && MPI. Initialized () && ucomm. val != - 1
42- instance. comm = ucomm
43- elseif isdefined (:MPI ) && MPI. Finalized ()
44- error (" MPI is already finalized!" )
45- else
46- # doing nothing
47- end
48- end
49-
50- return instance
51- end
52-
53- end
54- function freeMPIWrapper (instance:: MPIWrapper )
55- if isdefined (:MPI ) && MPI. Initialized () && ! MPI. Finalized ()
56- MPI. Finalize ()
57- end
58- end
59-
60- const mpiWrapper = MPIWrapper ();
61-
62- type StructureData
63- probability:: Dict{Int,Float64}
64- children:: Dict{Int,JuMP.Model}
65- parent
66- num_scen:: Int
67- othermap:: Dict{JuMP.Variable,JuMP.Variable}
68- mpiWrapper:: MPIWrapper
69- end
70- else
71- type StructureData
72- probability:: Dict{Int,Float64}
73- children:: Dict{Int,JuMP.Model}
74- parent
75- num_scen:: Int
76- othermap:: Dict{JuMP.Variable,JuMP.Variable}
77- end
26+ type StructureData
27+ probability:: Dict{Int,Float64}
28+ children:: Dict{Int,JuMP.Model}
29+ parent
30+ num_scen:: Int
31+ othermap:: Dict{JuMP.Variable,JuMP.Variable}
32+ MPIWrapper # Empty unless StructJuMPwithMPI fills it
7833end
7934default_probability (m:: JuMP.Model ) = 1 / num_scenarios (m)
8035default_probability (:: Void ) = 1.0
@@ -96,14 +51,21 @@ function structprinthook(io::IO, m::Model)
9651 end
9752end
9853
54+ type DummyMPIWrapper
55+ comm:: Int
56+ init:: Function
57+
58+ DummyMPIWrapper () = new (- 1 ,identity)
59+ end
60+ const dummy_mpi_wrapper = DummyMPIWrapper ()
61+
9962# Constructor with the number of scenarios
100- function StructuredModel (;solver= JuMP. UnsetSolver (), parent= nothing , same_children_as= nothing , id= 0 , comm= isdefined (:MPI ) ? MPI. Comm (- 1 ) : - 1 , num_scenarios:: Int = 0 , prob:: Float64 = default_probability (parent))
63+ function StructuredModel (;solver= JuMP. UnsetSolver (), parent= nothing , same_children_as= nothing , id= 0 , comm= nothing , num_scenarios:: Int = 0 , prob:: Float64 = default_probability (parent), mpi_wrapper= dummy_mpi_wrapper)
64+ _comm = (comm == nothing ? mpi_wrapper. comm : comm)
10165 m = JuMP. Model (solver= solver)
10266 if parent === nothing
10367 id = 0
104- if isdefined (:MPI )
105- mpiWrapper. init (comm)
106- end
68+ mpi_wrapper. init (_comm)
10769 if isdefined (:StructJuMPSolverInterface )
10870 JuMP. setsolvehook (m,StructJuMPSolverInterface. sj_solve)
10971 end
@@ -115,20 +77,16 @@ function StructuredModel(;solver=JuMP.UnsetSolver(), parent=nothing, same_childr
11577 end
11678
11779 if same_children_as != = nothing
118- if ! isa (same_children_as, JuMP. Model) || ! haskey (same_children_as. ext, :Stochastic )
119- error (" The JuMP model given for the argument `same_children_as' is not valid. Please create it using the `StructuredModel' function." )
120- end
121- probability = same_children_as. ext[:Stochastic ]. probability
122- children = same_children_as. ext[:Stochastic ]. children
123- else
124- probability = Dict {Int, Float64} ()
125- children = Dict {Int, JuMP.Model} ()
126- end
127- if isdefined (:MPI )
128- m. ext[:Stochastic ] = StructureData (probability, children, parent, num_scenarios, Dict {JuMP.Variable,JuMP.Variable} (), mpiWrapper)
80+ if ! isa (same_children_as, JuMP. Model) || ! haskey (same_children_as. ext, :Stochastic )
81+ error (" The JuMP model given for the argument `same_children_as' is not valid. Please create it using the `StructuredModel' function." )
82+ end
83+ probability = same_children_as. ext[:Stochastic ]. probability
84+ children = same_children_as. ext[:Stochastic ]. children
12985 else
130- m. ext[:Stochastic ] = StructureData (probability, children, parent, num_scenarios, Dict {JuMP.Variable,JuMP.Variable} ())
86+ probability = Dict {Int, Float64} ()
87+ children = Dict {Int, JuMP.Model} ()
13188 end
89+ m. ext[:Stochastic ] = StructureData (probability, children, parent, num_scenarios, Dict {JuMP.Variable,JuMP.Variable} (), mpi_wrapper)
13290
13391 # Printing children is important as well
13492 JuMP. setprinthook (m, structprinthook)
@@ -146,37 +104,12 @@ getchildren(m::JuMP.Model) = getStructure(m).children::Dict{Int,JuMP.Model}
146104getprobability (m:: JuMP.Model ) = getStructure (m). probability:: Dict{Int, Float64}
147105num_scenarios (m:: JuMP.Model ) = getStructure (m). num_scen:: Int
148106
149-
150- function getMyRank ()
151- myrank = 0 ;
152- mysize = 1 ;
153- if isdefined (:MPI ) && MPI. Initialized () && ! MPI. Finalized ()
154- comm = MPI. COMM_WORLD
155- mysize = MPI. Comm_size (comm)
156- myrank = MPI. Comm_rank (comm)
157- end
158- return myrank,mysize
159- end
160-
161- function getProcIdxSet (numScens:: Integer )
162- mysize = 1 ;
163- myrank = 0 ;
164- if isdefined (:MPI ) == true && MPI. Initialized () == true
165- comm = MPI. COMM_WORLD
166- mysize = MPI. Comm_size (comm)
167- myrank = MPI. Comm_rank (comm)
168- end
169- # Why don't we just take a round-and-robin?
170- proc_idx_set = Int[];
171- for s = myrank: mysize: (numScens- 1 )
172- push! (proc_idx_set, s+ 1 );
173- end
174- return proc_idx_set;
175- end
107+ getProcIdxSet (dummy_mpi_wrapper:: DummyMPIWrapper , num_scenarios) = 1 : num_scenarios
176108
177109function getProcIdxSet (m:: JuMP.Model )
110+ haskey (m. ext[:Stochastic ]) || error (" Cannot use @second_stage without using the StructuredModel constructor" )
178111 numScens = num_scenarios (m)
179- return getProcIdxSet (numScens)
112+ return getProcIdxSet (getStructure (m) . mpi_wrapper, numScens)
180113end
181114
182115macro second_stage (m,ind,code)
0 commit comments