Skip to content

Commit c4560bf

Browse files
bknuevenbknueven
authored
Merge pull request #432 from bknueven/sensi_rho
SensiRho part 1
2 parents c03fac5 + bee096e commit c4560bf

File tree

6 files changed

+900
-1
lines changed

6 files changed

+900
-1
lines changed

Diff for: ‎mpisppy/extensions/sensi_rho.py

+243
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,243 @@
1+
###############################################################################
2+
# mpi-sppy: MPI-based Stochastic Programming in PYthon
3+
#
4+
# Copyright (c) 2024, Lawrence Livermore National Security, LLC, Alliance for
5+
# Sustainable Energy, LLC, The Regents of the University of California, et al.
6+
# All rights reserved. Please see the files COPYRIGHT.md and LICENSE.md for
7+
# full copyright and license information.
8+
###############################################################################
9+
10+
import numpy as np
11+
12+
import pyomo.environ as pyo
13+
from pyomo.contrib.pynumero.linalg.scipy_interface import ScipyLU
14+
15+
import mpisppy.extensions.extension
16+
import mpisppy.MPI as MPI
17+
from mpisppy.utils.kkt.interface import InteriorPointInterface
18+
19+
20+
class SensiRho(mpisppy.extensions.extension.Extension):
21+
"""
22+
Rho determination algorithm using nonant sensitivities
23+
"""
24+
25+
def __init__(self, ph):
26+
self.ph = ph
27+
28+
self.multiplier = 1.0
29+
30+
if (
31+
"sensi_rho_options" in ph.options
32+
and "multiplier" in ph.options["sensi_rho_options"]
33+
):
34+
self.multiplier = ph.options["sensi_rho_options"]["multiplier"]
35+
36+
@staticmethod
37+
def _compute_rho_min_max(ph, npop, mpiop, start):
38+
local_nodenames = []
39+
local_xmaxmin = {}
40+
global_xmaxmin = {}
41+
42+
for k, s in ph.local_scenarios.items():
43+
nlens = s._mpisppy_data.nlens
44+
for node in s._mpisppy_node_list:
45+
if node.name not in local_nodenames:
46+
ndn = node.name
47+
local_nodenames.append(ndn)
48+
nlen = nlens[ndn]
49+
50+
local_xmaxmin[ndn] = start * np.ones(nlen, dtype="d")
51+
global_xmaxmin[ndn] = np.zeros(nlen, dtype="d")
52+
53+
for k, s in ph.local_scenarios.items():
54+
nlens = s._mpisppy_data.nlens
55+
rho = s._mpisppy_model.rho
56+
for node in s._mpisppy_node_list:
57+
ndn = node.name
58+
xmaxmin = local_xmaxmin[ndn]
59+
60+
xmaxmin_partial = np.fromiter(
61+
(rho[ndn,i]._value for i, _ in enumerate(node.nonant_vardata_list)),
62+
dtype="d",
63+
count=nlens[ndn],
64+
)
65+
xmaxmin = npop(xmaxmin, xmaxmin_partial)
66+
local_xmaxmin[ndn] = xmaxmin
67+
68+
for nodename in local_nodenames:
69+
ph.comms[nodename].Allreduce(
70+
[local_xmaxmin[nodename], MPI.DOUBLE],
71+
[global_xmaxmin[nodename], MPI.DOUBLE],
72+
op=mpiop,
73+
)
74+
75+
xmaxmin_dict = {}
76+
for ndn, global_xmaxmin_dict in global_xmaxmin.items():
77+
for i, v in enumerate(global_xmaxmin_dict):
78+
xmaxmin_dict[ndn, i] = v
79+
80+
return xmaxmin_dict
81+
82+
@staticmethod
83+
def _compute_rho_avg(ph):
84+
local_nodenames = []
85+
local_avg = {}
86+
global_avg = {}
87+
88+
for k, s in ph.local_scenarios.items():
89+
nlens = s._mpisppy_data.nlens
90+
rho = s._mpisppy_model.rho
91+
for node in s._mpisppy_node_list:
92+
if node.name not in local_nodenames:
93+
ndn = node.name
94+
local_nodenames.append(ndn)
95+
nlen = nlens[ndn]
96+
97+
local_avg[ndn] = np.zeros(nlen, dtype="d")
98+
global_avg[ndn] = np.zeros(nlen, dtype="d")
99+
100+
for k, s in ph.local_scenarios.items():
101+
nlens = s._mpisppy_data.nlens
102+
rho = s._mpisppy_model.rho
103+
for node in s._mpisppy_node_list:
104+
ndn = node.name
105+
106+
local_rhos = np.fromiter(
107+
(rho[ndn,i]._value for i, _ in enumerate(node.nonant_vardata_list)),
108+
dtype="d",
109+
count=nlens[ndn],
110+
)
111+
# print(f"{k=}, {local_rhos=}, {s._mpisppy_probability=}, {s._mpisppy_data.prob_coeff[ndn]=}")
112+
# TODO: is this the right thing, or should it be s._mpisppy_probability?
113+
local_rhos *= s._mpisppy_data.prob_coeff[ndn]
114+
115+
local_avg[ndn] += local_rhos
116+
117+
for nodename in local_nodenames:
118+
ph.comms[nodename].Allreduce(
119+
[local_avg[nodename], MPI.DOUBLE],
120+
[global_avg[nodename], MPI.DOUBLE],
121+
op=MPI.SUM,
122+
)
123+
124+
rhoavg_dict = {}
125+
for ndn, global_rhoavg_dict in global_avg.items():
126+
for i, v in enumerate(global_rhoavg_dict):
127+
rhoavg_dict[ndn, i] = v
128+
129+
return rhoavg_dict
130+
131+
@staticmethod
132+
def _compute_rho_max(ph):
133+
return SensiRho._compute_rho_min_max(ph, np.maximum, MPI.MAX, -np.inf)
134+
135+
@staticmethod
136+
def _compute_rho_min(ph):
137+
return SensiRho._compute_rho_min_max(ph, np.minimum, MPI.MIN, np.inf)
138+
139+
def pre_iter0(self):
140+
pass
141+
142+
def post_iter0(self):
143+
ph = self.ph
144+
145+
# first, solve the subproblems with Ipopt,
146+
# and gather sensitivity information
147+
ipopt = pyo.SolverFactory("ipopt")
148+
nonant_sensis = {}
149+
for k, s in ph.local_subproblems.items():
150+
solution_cache = pyo.ComponentMap()
151+
for var in s.component_data_objects(pyo.Var):
152+
solution_cache[var] = var._value
153+
relax_int = pyo.TransformationFactory('core.relax_integer_vars')
154+
relax_int.apply_to(s)
155+
156+
assert hasattr(s, "_relaxed_integer_vars")
157+
158+
# add the needed suffixes / remove later
159+
s.ipopt_zL_out = pyo.Suffix(direction=pyo.Suffix.IMPORT)
160+
s.ipopt_zU_out = pyo.Suffix(direction=pyo.Suffix.IMPORT)
161+
s.dual = pyo.Suffix(direction=pyo.Suffix.IMPORT_EXPORT)
162+
163+
results = ipopt.solve(s)
164+
pyo.assert_optimal_termination(results)
165+
166+
kkt_builder = InteriorPointInterface(s)
167+
kkt_builder.set_barrier_parameter(1e-9)
168+
kkt_builder.set_bounds_relaxation_factor(1e-8)
169+
#rhs = kkt_builder.evaluate_primal_dual_kkt_rhs()
170+
#print(f"{rhs}")
171+
#print(f"{rhs.flatten()}")
172+
kkt = kkt_builder.evaluate_primal_dual_kkt_matrix()
173+
174+
# print(f"{kkt=}")
175+
# could do better than SuperLU
176+
kkt_lu = ScipyLU()
177+
# always regularize equality constraints
178+
kkt_builder.regularize_equality_gradient(kkt=kkt, coef=-1e-8, copy_kkt=False)
179+
kkt_lu.do_numeric_factorization(kkt, raise_on_error=True)
180+
181+
grad_vec = np.zeros(kkt.shape[1])
182+
grad_vec[0:kkt_builder._nlp.n_primals()] = kkt_builder._nlp.evaluate_grad_objective()
183+
184+
grad_vec_kkt_inv = kkt_lu._lu.solve(grad_vec, "T")
185+
186+
for scenario_name in s.scen_list:
187+
nonant_sensis[scenario_name] = {}
188+
rho = ph.local_scenarios[scenario_name]._mpisppy_model.rho
189+
for ndn_i, v in ph.local_scenarios[scenario_name]._mpisppy_data.nonant_indices.items():
190+
var_idx = kkt_builder._nlp._vardata_to_idx[v]
191+
192+
y_vec = np.zeros(kkt.shape[0])
193+
y_vec[var_idx] = 1.0
194+
195+
x_denom = y_vec.T @ kkt_lu._lu.solve(y_vec)
196+
x = (-1 / x_denom)
197+
e_x = x * y_vec
198+
199+
sensitivity = grad_vec_kkt_inv @ -e_x
200+
# print(f"df/d{v.name}: {sensitivity:.2e}, ∂f/∂{v.name}: {grad_vec[var_idx]:.2e}, "
201+
# f"rho {v.name}: {ph.local_scenarios[scenario_name]._mpisppy_model.rho[ndn_i]._value:.2e}, ",
202+
# f"value: {v._value:.2e}"
203+
# )
204+
205+
rho[ndn_i]._value = abs(sensitivity)
206+
207+
relax_int.apply_to(s, options={"undo":True})
208+
assert not hasattr(s, "_relaxed_integer_vars")
209+
del s.ipopt_zL_out
210+
del s.ipopt_zU_out
211+
del s.dual
212+
for var, val in solution_cache.items():
213+
var._value = val
214+
215+
for s in ph.local_scenarios.values():
216+
xbars = s._mpisppy_model.xbars
217+
for ndn_i, rho in s._mpisppy_model.rho.items():
218+
nv = s._mpisppy_data.nonant_indices[ndn_i] # var_data object
219+
rho._value = rho._value / max(1, abs(nv._value - xbars[ndn_i]._value))
220+
rho._value *= self.multiplier
221+
# if ph.cylinder_rank == 0:
222+
# print(f"{s.name=}, {nv.name=}, {rho.value=}")
223+
224+
rhoavg = self._compute_rho_avg(ph)
225+
for s in ph.local_scenarios.values():
226+
xbars = s._mpisppy_model.xbars
227+
for ndn_i, rho in s._mpisppy_model.rho.items():
228+
rho._value = rhoavg[ndn_i]
229+
# if ph.cylinder_rank == 0:
230+
# nv = s._mpisppy_data.nonant_indices[ndn_i] # var_data object
231+
# print(f"{s.name=}, {nv.name=}, {rho.value=}")
232+
233+
if ph.cylinder_rank == 0:
234+
print("Rho values updated by SensiRho Extension")
235+
236+
def miditer(self):
237+
pass
238+
239+
def enditer(self):
240+
pass
241+
242+
def post_everything(self):
243+
pass

Diff for: ‎mpisppy/generic_cylinders.py

+10-1
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,7 @@ def _parse_args(m):
7272
cfg.reduced_costs_args()
7373
cfg.sep_rho_args()
7474
cfg.coeff_rho_args()
75+
cfg.sensi_rho_args()
7576

7677
cfg.parse_command_line(f"mpi-sppy for {cfg.module_name}")
7778
return cfg
@@ -99,7 +100,7 @@ def _name_lists(module, cfg):
99100
def _do_decomp(module, cfg, scenario_creator, scenario_creator_kwargs, scenario_denouement):
100101
rho_setter = module._rho_setter if hasattr(module, '_rho_setter') else None
101102
if cfg.default_rho is None and rho_setter is None:
102-
if cfg.sep_rho or cfg.coeff_rho:
103+
if cfg.sep_rho or cfg.coeff_rho or cfg.sensi_rho:
103104
cfg.default_rho = 1
104105
else:
105106
raise RuntimeError("No rho_setter so a default must be specified via --default-rho")
@@ -167,6 +168,9 @@ def _do_decomp(module, cfg, scenario_creator, scenario_creator_kwargs, scenario_
167168

168169
if cfg.coeff_rho:
169170
vanilla.add_coeff_rho(hub_dict, cfg)
171+
172+
if cfg.sensi_rho:
173+
vanilla.add_sensi_rho(hub_dict, cfg)
170174

171175
if len(ext_classes) != 0:
172176
hub_dict['opt_kwargs']['extensions'] = MultiExtension
@@ -215,6 +219,9 @@ def _do_decomp(module, cfg, scenario_creator, scenario_creator_kwargs, scenario_
215219
vanilla.add_sep_rho(ph_ob_spoke, cfg)
216220
if cfg.coeff_rho:
217221
vanilla.add_coeff_rho(ph_ob_spoke, cfg)
222+
if cfg.sensi_rho:
223+
vanilla.add_sensi_rho(ph_ob_spoke, cfg)
224+
218225

219226
# subgradient outer bound spoke
220227
if cfg.subgradient:
@@ -227,6 +234,8 @@ def _do_decomp(module, cfg, scenario_creator, scenario_creator_kwargs, scenario_
227234
vanilla.add_sep_rho(subgradient_spoke, cfg)
228235
if cfg.coeff_rho:
229236
vanilla.add_coeff_rho(subgradient_spoke, cfg)
237+
if cfg.sensi_rho:
238+
vanilla.add_sensi_rho(subgradient_spoke, cfg)
230239

231240
# xhat shuffle bound spoke
232241
if cfg.xhatshuffle:

Diff for: ‎mpisppy/utils/cfg_vanilla.py

+5
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@
4242
from mpisppy.extensions.reduced_costs_fixer import ReducedCostsFixer
4343
from mpisppy.extensions.sep_rho import SepRho
4444
from mpisppy.extensions.coeff_rho import CoeffRho
45+
from mpisppy.extensions.sensi_rho import SensiRho
4546
from mpisppy.utils.wxbarreader import WXBarReader
4647
from mpisppy.utils.wxbarwriter import WXBarWriter
4748

@@ -209,6 +210,10 @@ def add_coeff_rho(hub_dict, cfg):
209210
hub_dict = extension_adder(hub_dict,CoeffRho)
210211
hub_dict["opt_kwargs"]["options"]["coeff_rho_options"] = {"multiplier" : cfg.coeff_rho_multiplier}
211212

213+
def add_sensi_rho(hub_dict, cfg):
214+
hub_dict = extension_adder(hub_dict,SensiRho)
215+
hub_dict["opt_kwargs"]["options"]["sensi_rho_options"] = {"multiplier" : cfg.sensi_rho_multiplier}
216+
212217
def add_cross_scenario_cuts(hub_dict,
213218
cfg,
214219
):

Diff for: ‎mpisppy/utils/config.py

+11
Original file line numberDiff line numberDiff line change
@@ -421,6 +421,17 @@ def sep_rho_args(self):
421421
default=1.0)
422422

423423

424+
def sensi_rho_args(self):
425+
self.add_to_config("sensi_rho",
426+
description="have a SensiRho extension",
427+
domain=bool,
428+
default=False)
429+
self.add_to_config("sensi_rho_multiplier",
430+
description="multiplier for SensiRho (default 1.0)",
431+
domain=float,
432+
default=1.0)
433+
434+
424435
def coeff_rho_args(self):
425436
self.add_to_config("coeff_rho",
426437
description="have a CoeffRho extension",

0 commit comments

Comments
 (0)