def test_linear_solution_cache(self):
# Test derivatives across a converged Sellar model. When caching
# is performed, the second solve takes less iterations than the
# first one.
# Forward mode
prob = om.Problem()
model = prob.model
model.add_subsystem('px', om.IndepVarComp('x', 1.0), promotes=['x'])
model.add_subsystem('d1', Comp4LinearCacheTest(), promotes=['x', 'y'])
model.nonlinear_solver = om.NonlinearBlockGS()
model.linear_solver = om.PETScKrylov()
model.add_design_var('x', cache_linear_solution=True)
model.add_objective('y', cache_linear_solution=True)
prob.setup(mode='fwd')
prob.set_solver_print(level=0)
prob.run_model()
J = prob.driver._compute_totals(of=['y'], wrt=['x'], use_abs_names=False,
return_format='flat_dict')
icount1 = prob.model.linear_solver._iter_count
J = prob.driver._compute_totals(of=['y'], wrt=['x'], use_abs_names=False,
return_format='flat_dict')
icount2 = prob.model.linear_solver._iter_count
# Should take less iterations when starting from previous solution.
self.assertTrue(icount2 < icount1)
# Reverse mode
prob = om.Problem()
model = prob.model
model.add_subsystem('px', om.IndepVarComp('x', 1.0), promotes=['x'])
model.add_subsystem('d1', Comp4LinearCacheTest(), promotes=['x', 'y'])
model.nonlinear_solver = om.NonlinearBlockGS()
model.linear_solver = om.PETScKrylov()
model.add_design_var('x', cache_linear_solution=True)
model.add_objective('y', cache_linear_solution=True)
prob.setup(mode='rev')
prob.set_solver_print(level=0)
prob.run_model()
J = prob.driver._compute_totals(of=['y'], wrt=['x'], use_abs_names=False,
return_format='flat_dict')
icount1 = prob.model.linear_solver._iter_count
J = prob.driver._compute_totals(of=['y'], wrt=['x'], use_abs_names=False,
return_format='flat_dict')
icount2 = prob.model.linear_solver._iter_count
# Should take less iterations when starting from previous solution.
self.assertTrue(icount2 < icount1)
def test_options(self):
"""Verify that the PETScKrylov specific options are declared."""
group = om.Group()
group.linear_solver = om.PETScKrylov()
assert (group.linear_solver.options['ksp_type'] == 'fgmres')
def test_specify_precon(self):
import numpy as np
import openmdao.api as om
from openmdao.test_suite.components.sellar import SellarDis1withDerivatives, \
SellarDis2withDerivatives
prob = om.Problem()
model = prob.model
model.add_subsystem('d1', SellarDis1withDerivatives(), promotes=['x', 'z', 'y1', 'y2'])
model.add_subsystem('d2', SellarDis2withDerivatives(), promotes=['z', 'y1', 'y2'])
model.add_subsystem('obj_cmp', om.ExecComp('obj = x**2 + z[1] + y1 + exp(-y2)',
z=np.array([0.0, 0.0]), x=0.0),
promotes=['obj', 'x', 'z', 'y1', 'y2'])
model.add_subsystem('con_cmp1', om.ExecComp('con1 = 3.16 - y1'), promotes=['con1', 'y1'])
model.add_subsystem('con_cmp2', om.ExecComp('con2 = y2 - 24.0'), promotes=['con2', 'y2'])
model.nonlinear_solver = om.NewtonSolver(solve_subsystems=False)
model.linear_solver = om.PETScKrylov()
model.linear_solver.precon = om.LinearBlockGS()
model.linear_solver.precon.options['maxiter'] = 2
prob.setup()
prob.set_val('x', 1.)
prob.set_val('z', np.array([5.0, 2.0]))
prob.run_model()
assert_near_equal(prob.get_val('y1'), 25.58830273, .00001)
assert_near_equal(prob.get_val('y2'), 12.05848819, .00001)
def test_error_under_cs(self):
"""Verify that PETScKrylov abides by the 'maxiter' option."""
prob = om.Problem()
model = prob.model
model.add_subsystem('px', om.IndepVarComp('x', 1.0), promotes=['x'])
model.add_subsystem('pz', om.IndepVarComp('z', np.array([5.0, 2.0])), promotes=['z'])
model.add_subsystem('d1', SellarDis1withDerivatives(), promotes=['x', 'z', 'y1', 'y2'])
model.add_subsystem('d2', SellarDis2withDerivatives(), promotes=['z', 'y1', 'y2'])
model.add_subsystem('obj_cmp', om.ExecComp('obj = x**2 + z[1] + y1 + exp(-y2)',
z=np.array([0.0, 0.0]), x=0.0),
promotes=['obj', 'x', 'z', 'y1', 'y2'])
model.add_subsystem('con_cmp1', om.ExecComp('con1 = 3.16 - y1'), promotes=['con1', 'y1'])
model.add_subsystem('con_cmp2', om.ExecComp('con2 = y2 - 24.0'), promotes=['con2', 'y2'])
model.nonlinear_solver = om.NewtonSolver(solve_subsystems=False)
model.linear_solver = om.PETScKrylov()
model.approx_totals(method='cs')
prob.setup(mode='fwd')
prob.set_solver_print(level=0)
prob.run_model()
with self.assertRaises(RuntimeError) as cm:
J = prob.compute_totals(of=['obj'], wrt=['z'])
msg = 'PETScKrylov in <model> <class Group>: PETScKrylov solver is not supported under complex step.'
self.assertEqual(str(cm.exception), msg)
def test_specify_precon_left(self):
import numpy as np
import openmdao.api as om
from openmdao.test_suite.components.sellar import SellarDis1withDerivatives, \
SellarDis2withDerivatives
prob = om.Problem()
model = prob.model
model.add_subsystem('px', om.IndepVarComp('x', 1.0), promotes=['x'])
model.add_subsystem('pz', om.IndepVarComp('z', np.array([5.0, 2.0])), promotes=['z'])
model.add_subsystem('d1', SellarDis1withDerivatives(), promotes=['x', 'z', 'y1', 'y2'])
model.add_subsystem('d2', SellarDis2withDerivatives(), promotes=['z', 'y1', 'y2'])
model.add_subsystem('obj_cmp', om.ExecComp('obj = x**2 + z[1] + y1 + exp(-y2)',
z=np.array([0.0, 0.0]), x=0.0),
promotes=['obj', 'x', 'z', 'y1', 'y2'])
model.add_subsystem('con_cmp1', om.ExecComp('con1 = 3.16 - y1'), promotes=['con1', 'y1'])
model.add_subsystem('con_cmp2', om.ExecComp('con2 = y2 - 24.0'), promotes=['con2', 'y2'])
model.nonlinear_solver = om.NewtonSolver()
model.linear_solver = om.PETScKrylov()
model.linear_solver.precon = om.DirectSolver()
model.linear_solver.options['precon_side'] = 'left'
model.linear_solver.options['ksp_type'] = 'richardson'
prob.setup()
prob.run_model()
assert_rel_error(self, prob['y1'], 25.58830273, .00001)
assert_rel_error(self, prob['y2'], 12.05848819, .00001)
def test_specify_solver(self):
import numpy as np
import openmdao.api as om
from openmdao.test_suite.components.sellar import SellarDis1withDerivatives, SellarDis2withDerivatives
prob = om.Problem()
model = prob.model
model.add_subsystem('px', om.IndepVarComp('x', 1.0), promotes=['x'])
model.add_subsystem('pz', om.IndepVarComp('z', np.array([5.0, 2.0])), promotes=['z'])
model.add_subsystem('d1', SellarDis1withDerivatives(), promotes=['x', 'z', 'y1', 'y2'])
model.add_subsystem('d2', SellarDis2withDerivatives(), promotes=['z', 'y1', 'y2'])
model.add_subsystem('obj_cmp', om.ExecComp('obj = x**2 + z[1] + y1 + exp(-y2)',
z=np.array([0.0, 0.0]), x=0.0),
promotes=['obj', 'x', 'z', 'y1', 'y2'])
model.add_subsystem('con_cmp1', om.ExecComp('con1 = 3.16 - y1'), promotes=['con1', 'y1'])
model.add_subsystem('con_cmp2', om.ExecComp('con2 = y2 - 24.0'), promotes=['con2', 'y2'])
model.nonlinear_solver = om.NonlinearBlockGS()
model.linear_solver = om.PETScKrylov()
prob.setup()
prob.run_model()
wrt = ['z']
of = ['obj']
J = prob.compute_totals(of=of, wrt=wrt, return_format='flat_dict')
assert_rel_error(self, J['obj', 'z'][0][0], 9.61001056, .00001)
assert_rel_error(self, J['obj', 'z'][0][1], 1.78448534, .00001)
def setup(self):
self.add_input('a', val=10., units='m')
rank = self.comm.rank
GLOBAL_SIZE = 15
sizes, offsets = evenly_distrib_idxs(self.comm.size, GLOBAL_SIZE)
self.add_output('states', shape=int(sizes[rank]))
self.add_output('out_var', shape=1)
self.local_size = sizes[rank]
self.linear_solver = om.PETScKrylov()
self.linear_solver.precon = om.LinearUserDefined(solve_function=self.mysolve)
def test_feature_maxiter(self):
import numpy as np
import openmdao.api as om
from openmdao.test_suite.components.sellar import SellarDis1withDerivatives, SellarDis2withDerivatives
prob = om.Problem()
model = prob.model
model.add_subsystem('d1',
SellarDis1withDerivatives(),
promotes=['x', 'z', 'y1', 'y2'])
model.add_subsystem('d2',
SellarDis2withDerivatives(),
promotes=['z', 'y1', 'y2'])
model.add_subsystem('obj_cmp',
om.ExecComp('obj = x**2 + z[1] + y1 + exp(-y2)',
z=np.array([0.0, 0.0]),
x=0.0),
promotes=['obj', 'x', 'z', 'y1', 'y2'])
model.add_subsystem('con_cmp1',
om.ExecComp('con1 = 3.16 - y1'),
promotes=['con1', 'y1'])
model.add_subsystem('con_cmp2',
om.ExecComp('con2 = y2 - 24.0'),
promotes=['con2', 'y2'])
model.nonlinear_solver = om.NonlinearBlockGS()
model.linear_solver = om.PETScKrylov()
model.linear_solver.options['maxiter'] = 3
prob.setup()
prob.set_val('x', 1.)
prob.set_val('z', np.array([5.0, 2.0]))
prob.run_model()
wrt = ['z']
of = ['obj']
J = prob.compute_totals(of=of, wrt=wrt, return_format='flat_dict')
assert_near_equal(J['obj', 'z'][0][0], 4.93218027, .00001)
assert_near_equal(J['obj', 'z'][0][1], 1.73406455, .00001)
def setup(self):
self.options['distributed'] = True
self.add_input('a', val=10., units='m', src_indices=[0])
rank = self.comm.rank
GLOBAL_SIZE = 5
sizes, offsets = evenly_distrib_idxs(self.comm.size, GLOBAL_SIZE)
self.add_output('states', shape=int(sizes[rank]))
self.add_output('out_var', shape=1)
self.local_size = sizes[rank]
self.linear_solver = om.PETScKrylov()
def test_method(self):
p = om.Problem()
p.model.add_subsystem('des_vars', om.IndepVarComp('a', val=10., units='m'), promotes=['*'])
p.model.add_subsystem('icomp', DistribStateImplicit(), promotes=['*'])
model = p.model
model.linear_solver = om.PETScKrylov()
model.linear_solver.precon = om.LinearRunOnce()
p.setup(mode='rev', check=False)
p.run_model()
jac = p.compute_totals(of=['out_var'], wrt=['a'], return_format='dict')
assert_near_equal(15.0, jac['out_var']['a'][0][0])
def test_feature_rtol(self):
prob = om.Problem()
model = prob.model
model.add_subsystem('d1',
SellarDis1withDerivatives(),
promotes=['x', 'z', 'y1', 'y2'])
model.add_subsystem('d2',
SellarDis2withDerivatives(),
promotes=['z', 'y1', 'y2'])
model.add_subsystem('obj_cmp',
om.ExecComp('obj = x**2 + z[1] + y1 + exp(-y2)',
z=np.array([0.0, 0.0]),
x=0.0),
promotes=['obj', 'x', 'z', 'y1', 'y2'])
model.add_subsystem('con_cmp1',
om.ExecComp('con1 = 3.16 - y1'),
promotes=['con1', 'y1'])
model.add_subsystem('con_cmp2',
om.ExecComp('con2 = y2 - 24.0'),
promotes=['con2', 'y2'])
model.nonlinear_solver = om.NonlinearBlockGS()
model.linear_solver = om.PETScKrylov()
model.linear_solver.options['rtol'] = 1.0e-20
prob.setup()
prob.set_val('x', 1.)
prob.set_val('z', np.array([5.0, 2.0]))
prob.run_model()
wrt = ['z']
of = ['obj']
J = prob.compute_totals(of=of, wrt=wrt, return_format='flat_dict')
assert_near_equal(J['obj', 'z'][0][0], 9.61001055699, .00001)
assert_near_equal(J['obj', 'z'][0][1], 1.78448533563, .00001)
def test_hierarchy_iprint(self):
prob = om.Problem()
model = prob.model
model.add_subsystem('pz', om.IndepVarComp('z', np.array([5.0, 2.0])))
sub1 = model.add_subsystem('sub1', om.Group())
sub2 = sub1.add_subsystem('sub2', om.Group())
g1 = sub2.add_subsystem('g1', SubSellar())
g2 = model.add_subsystem('g2', SubSellar())
model.connect('pz.z', 'sub1.sub2.g1.z')
model.connect('sub1.sub2.g1.y2', 'g2.x')
model.connect('g2.y2', 'sub1.sub2.g1.x')
model.nonlinear_solver = om.NewtonSolver()
model.linear_solver = om.LinearBlockGS()
model.nonlinear_solver.options['solve_subsystems'] = True
model.nonlinear_solver.options['max_sub_solves'] = 0
g1.nonlinear_solver = om.NewtonSolver(solve_subsystems=False)
g1.linear_solver = om.LinearBlockGS()
g2.nonlinear_solver = om.NewtonSolver(solve_subsystems=False)
g2.linear_solver = om.PETScKrylov()
g2.linear_solver.precon = om.LinearBlockGS()
g2.linear_solver.precon.options['maxiter'] = 2
prob.set_solver_print(level=2)
prob.setup()
# Conclude setup but don't run model.
prob.final_setup()
# if USE_PROC_FILES is not set, solver convergence messages
# should only appear on proc 0
output = run_model(prob)
if model.comm.rank == 0 or os.environ.get('USE_PROC_FILES'):
self.assertTrue(output.count('\nNL: Newton Converged') == 1)
else:
self.assertTrue(output.count('\nNL: Newton Converged') == 0)
def test_specify_precon_left(self):
prob = om.Problem()
model = prob.model
model.add_subsystem('d1',
SellarDis1withDerivatives(),
promotes=['x', 'z', 'y1', 'y2'])
model.add_subsystem('d2',
SellarDis2withDerivatives(),
promotes=['z', 'y1', 'y2'])
model.add_subsystem('obj_cmp',
om.ExecComp('obj = x**2 + z[1] + y1 + exp(-y2)',
z=np.array([0.0, 0.0]),
x=0.0),
promotes=['obj', 'x', 'z', 'y1', 'y2'])
model.add_subsystem('con_cmp1',
om.ExecComp('con1 = 3.16 - y1'),
promotes=['con1', 'y1'])
model.add_subsystem('con_cmp2',
om.ExecComp('con2 = y2 - 24.0'),
promotes=['con2', 'y2'])
model.nonlinear_solver = om.NewtonSolver(solve_subsystems=False)
model.linear_solver = om.PETScKrylov()
model.linear_solver.precon = om.DirectSolver()
model.linear_solver.options['precon_side'] = 'left'
model.linear_solver.options['ksp_type'] = 'richardson'
prob.setup()
prob.set_val('x', 1.)
prob.set_val('z', np.array([5.0, 2.0]))
prob.run_model()
assert_near_equal(prob.get_val('y1'), 25.58830273, .00001)
assert_near_equal(prob.get_val('y2'), 12.05848819, .00001)
def test_feature(self):
class CustomSolveImplicit(om.ImplicitComponent):
def setup(self):
self.add_input('a', val=10., units='m')
rank = self.comm.rank
GLOBAL_SIZE = 15
sizes, offsets = evenly_distrib_idxs(self.comm.size, GLOBAL_SIZE)
self.add_output('states', shape=int(sizes[rank]))
self.add_output('out_var', shape=1)
self.local_size = sizes[rank]
self.linear_solver = om.PETScKrylov()
self.linear_solver.precon = om.LinearUserDefined(solve_function=self.mysolve)
def solve_nonlinear(self, i, o):
o['states'] = i['a']
local_sum = np.zeros(1)
local_sum[0] = np.sum(o['states'])
tmp = np.zeros(1)
o['out_var'] = tmp[0]
def apply_nonlinear(self, i, o, r):
r['states'] = o['states'] - i['a']
local_sum = np.zeros(1)
local_sum[0] = np.sum(o['states'])
global_sum = np.zeros(1)
r['out_var'] = o['out_var'] - tmp[0]
def apply_linear(self, i, o, d_i, d_o, d_r, mode):
if mode == 'fwd':
if 'states' in d_o:
d_r['states'] += d_o['states']
local_sum = np.array([np.sum(d_o['states'])])
global_sum = np.zeros(1)
self.comm.Allreduce(local_sum, global_sum, op=MPI.SUM)
d_r['out_var'] -= global_sum
if 'out_var' in d_o:
d_r['out_var'] += d_o['out_var']
if 'a' in d_i:
d_r['states'] -= d_i['a']
elif mode == 'rev':
if 'states' in d_o:
d_o['states'] += d_r['states']
tmp = np.zeros(1)
if self.comm.rank == 0:
tmp[0] = d_r['out_var'].copy()
self.comm.Bcast(tmp, root=0)
d_o['states'] -= tmp
if 'out_var' in d_o:
d_o['out_var'] += d_r['out_var']
if 'a' in d_i:
d_i['a'] -= np.sum(d_r['states'])
def mysolve(self, d_outputs, d_residuals, mode):
r"""
Apply inverse jac product. The model is assumed to be in an unscaled state.
If mode is:
'fwd': d_residuals \|-> d_outputs
'rev': d_outputs \|-> d_residuals
Parameters
----------
d_outputs: Vector
Unscaled, dimensional quantities read via d_outputs[key].
d_residuals: Vector
Unscaled, dimensional quantities read via d_residuals[key].
mode: str
either 'fwd' or 'rev'
"""
# Note: we are just preconditioning with Identity as a proof of concept.
if mode == 'fwd':
d_outputs.set_vec(d_residuals)
elif mode == 'rev':
d_residuals.set_vec(d_outputs)
prob = om.Problem()
prob.model.add_subsystem('icomp', CustomSolveImplicit(), promotes=['*'])
prob.model.set_input_defaults('a', 10., units='m')
model = prob.model
model.linear_solver = om.PETScKrylov()
model.linear_solver.precon = om.LinearRunOnce()
prob.setup(mode='rev', check=False)
prob.run_model()
jac = prob.compute_totals(of=['out_var'], wrt=['a'], return_format='dict')
assert_near_equal(15.0, jac['out_var']['a'][0][0])
