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 = Problem()
model = prob.model
model.add_subsystem('px', IndepVarComp('x', 1.0), promotes=['x'])
model.add_subsystem('d1', Comp4LinearCacheTest(), promotes=['x', 'y'])
model.nonlinear_solver = NonlinearBlockGS()
model.linear_solver = 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'], global_names=False,
return_format='flat_dict')
icount1 = prob.model.linear_solver._iter_count
J = prob.driver._compute_totals(of=['y'], wrt=['x'], global_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 = Problem()
model = prob.model
model.add_subsystem('px', IndepVarComp('x', 1.0), promotes=['x'])
model.add_subsystem('d1', Comp4LinearCacheTest(), promotes=['x', 'y'])
model.nonlinear_solver = NonlinearBlockGS()
model.linear_solver = 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'], global_names=False,
return_format='flat_dict')
icount1 = prob.model.linear_solver._iter_count
J = prob.driver._compute_totals(of=['y'], wrt=['x'], global_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 = Group()
group.linear_solver = PETScKrylov()
assert (group.linear_solver.options['ksp_type'] == 'fgmres')
def test_specify_precon(self):
import numpy as np
from openmdao.api import Problem, Group, IndepVarComp, LinearBlockGS, PETScKrylov, \
NewtonSolver, ExecComp
from openmdao.test_suite.components.sellar import SellarDis1withDerivatives, \
SellarDis2withDerivatives
prob = Problem()
model = prob.model = Group()
model.add_subsystem('px', IndepVarComp('x', 1.0), promotes=['x'])
model.add_subsystem('pz', 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', 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', ExecComp('con1 = 3.16 - y1'), promotes=['con1', 'y1'])
model.add_subsystem('con_cmp2', ExecComp('con2 = y2 - 24.0'), promotes=['con2', 'y2'])
model.nonlinear_solver = NewtonSolver()
model.linear_solver = PETScKrylov()
model.linear_solver.precon = LinearBlockGS()
model.linear_solver.precon.options['maxiter'] = 2
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
from openmdao.api import Problem, Group, IndepVarComp, NonlinearBlockGS, PETScKrylov, ExecComp
from openmdao.test_suite.components.sellar import SellarDis1withDerivatives, SellarDis2withDerivatives
prob = Problem()
model = prob.model = Group()
model.add_subsystem('px', IndepVarComp('x', 1.0), promotes=['x'])
model.add_subsystem('pz', 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', 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', ExecComp('con1 = 3.16 - y1'), promotes=['con1', 'y1'])
model.add_subsystem('con_cmp2', ExecComp('con2 = y2 - 24.0'), promotes=['con2', 'y2'])
model.nonlinear_solver = NonlinearBlockGS()
model.linear_solver = 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 test_error_under_cs(self):
"""Verify that PETScKrylov abides by the 'maxiter' option."""
prob = Problem()
model = prob.model = Group()
model.add_subsystem('px', IndepVarComp('x', 1.0), promotes=['x'])
model.add_subsystem('pz', 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', 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', ExecComp('con1 = 3.16 - y1'), promotes=['con1', 'y1'])
model.add_subsystem('con_cmp2', ExecComp('con2 = y2 - 24.0'), promotes=['con2', 'y2'])
model.nonlinear_solver = NewtonSolver()
model.linear_solver = 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 solver is not supported under complex step.'
self.assertEqual(str(cm.exception), msg)
def setup(self):
ap = self.options['ap']
self.solver, self.mesh, self.dvgeo, self.dvcon = self.options['setup_cb'](self.comm)
# necessary to get some memory properly allocated
#self.solver.getResidual(ap)
des_vars, constraints = get_dvs_and_cons(ap, self.dvgeo, self.dvcon)
geo_vars, _ = get_dvs_and_cons(geo=self.dvgeo)
self.des_var_names = des_vars
self.geo_var_names = geo_vars
self.constraint_names = constraints
des_var_names = [dv[0][0] for dv in des_vars]
geo_var_names = [dv[0][0] for dv in geo_vars]
if self.options['owns_indeps']:
indeps = self.add_subsystem('indeps', IndepVarComp(), promotes=['*'])
for (args, kwargs) in des_vars:
name = args[0]
size = args[1]
value= kwargs['value']
if 'units' in kwargs:
indeps.add_output(name, value, units=kwargs['units'])
else:
indeps.add_output(name, value)
states = OM_STATES_COMP(ap=ap, dvgeo=self.dvgeo, solver=self.solver,
use_OM_KSP=self.options['use_OM_KSP'])
self.add_subsystem('states', states, promotes_inputs=des_var_names)
# this lets the OpenMDAO KSPsolver converge the adjoint using the ADflow PC
if self.options['use_OM_KSP']:
states.linear_solver = PETScKrylov(iprint=2, atol=1e-8, rtol=1e-8, maxiter=300, ksp_type='gmres')
states.linear_solver.precon = LinearUserDefined()
functionals = OM_FUNC_COMP(ap=ap, dvgeo=self.dvgeo, solver=self.solver)
self.add_subsystem('functionals', functionals,
promotes_inputs=des_var_names, max_procs=self.comm.size)
self.connect('states.states', 'functionals.states')
if self.dvcon is not None:
geocon = OM_GEOCON_COMP(dvgeo=self.dvgeo, dvcon=self.dvcon)
self.add_subsystem('geocon', geocon, promotes_inputs=geo_var_names)
for cons in self.dvcon.constraints.values():
for name, con in cons.items():
if self.comm.rank == 0:
print('adding nonlinear dvcon constarint: ', name)
self.add_constraint('geocon.%s' % name, lower=con.lower, upper=con.upper, ref=1/con.scale, vectorize_derivs=True)
for name, con in self.dvcon.linearCon.items():
if self.comm.rank == 0:
print('adding linear dvcon constraint: ', name)
self.add_constraint('geocon.%s' % name, lower=con.lower, upper=con.upper, linear=True, vectorize_derivs=True)
def setup(self):
self.add_input('a', val=10., units='m')
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 = PETScKrylov()
def setup(self):
self._mpi_proc_allocator.parallel = self.parallel
if self.parallel:
self.nonlinear_solver = NewtonSolver(solve_subsystems=False)
self.linear_solver = PETScKrylov()
else:
self.nonlinear_solver = NonlinearBlockGS()
self.linear_solver = LinearBlockGS()
self.add_subsystem('C1', ExecComp('z = 1 / 3. * y + x0'), promotes=['x0'])
self.add_subsystem('C2', ExecComp('z = 1 / 4. * y + x1'), promotes=['x1'])
self.connect('C1.z', 'C2.y')
self.connect('C2.z', 'C1.y')
self.parallel = not self.parallel
def test_method(self):
p = Problem()
p.model.add_subsystem('des_vars', IndepVarComp('a', val=10., units='m'), promotes=['*'])
p.model.add_subsystem('icomp', DistribStateImplicit(), promotes=['*'])
model = p.model
model.linear_solver = PETScKrylov()
model.linear_solver.precon = LinearRunOnce()
p.setup(vector_class=PETScVector, mode='rev', check=False)
p.run_model()
jac = p.compute_totals(of=['out_var'], wrt=['a'], return_format='dict')
assert_rel_error(self, 15.0, jac['out_var']['a'][0][0])
"""
k: local stiffness matrix of element_on
"""
i = np_mesh[element_ID, 0, 3] + 1
j = np_mesh[element_ID, 1, 3] + 1
m = np_mesh[element_ID, 2, 3] + 1
n = np_mesh[element_ID, 3, 3] + 1
k = np.matmul(np.matmul(V * np.transpose(B),D), B)
_K(K, k, i,j,m,n, element_ID)
from openmdao.api import Problem, Group, ImplicitComponent, IndepVarComp, NonlinearBlockGS, PETScKrylov
prob = Problem()
model = prob.model = Group()
np.random.seed(0)
num_elements = 400
model.add_subsystem('pK', IndepVarComp('K', np.random.rand(num_elements,num_elements)))
model.add_subsystem('pf', IndepVarComp('f', np.eye(num_elements,1)))
model.add_subsystem('Kdeqf', FEM3dComp(num_elements=num_elements,num_nodes=num_nodes))
model.connect('pK.K', 'Kdeqf.K')
model.connect('pf.f', 'Kdeqf.f')
model.nonlinear_solver = NonlinearBlockGS()
model.linear_solver = PETScKrylov()
prob.setup()
prob.run_model()
prob['Kdeqf.d']
def test_feature(self):
import numpy as np
from openmdao.api import Problem, ImplicitComponent, IndepVarComp, LinearRunOnce, PETScKrylov, PETScVector, LinearUserDefined
from openmdao.utils.array_utils import evenly_distrib_idxs
class CustomSolveImplicit(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 = PETScKrylov()
self.linear_solver.precon = 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'
Returns
-------
None or bool or (bool, float, float)
The bool is the failure flag; and the two floats are absolute and relative error.
"""
# 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)
return False, 0., 0.
prob = Problem()
prob.model.add_subsystem('des_vars',
IndepVarComp('a', val=10., units='m'),
promotes=['*'])
prob.model.add_subsystem('icomp',
CustomSolveImplicit(),
promotes=['*'])
model = prob.model
model.linear_solver = PETScKrylov()
model.linear_solver.precon = LinearRunOnce()
prob.setup(mode='rev', check=False)
prob.run_model()
jac = prob.compute_totals(of=['out_var'],
wrt=['a'],
return_format='dict')
assert_rel_error(self, 15.0, jac['out_var']['a'][0][0])