def test_coloring1(self):
mat, inshapes, outshapes = mat_factory(3, 2)
outsizes = [np.prod(shp) for shp in outshapes]
def func(a, b, c):
ivec = np.hstack([a.flat, b.flat, c.flat])
ovec = mat.dot(ivec)
x, y = ovec2outs(ovec, outsizes)
return x, y
f = (omf.wrap(func)
.add_inputs(a={'shape': inshapes[0]}, b={'shape': inshapes[1]}, c={'shape': inshapes[2]})
.add_outputs(x={'shape': outshapes[0]}, y={'shape': outshapes[1]})
.declare_coloring(wrt='*', method='cs', show_summary=False)
)
p = om.Problem()
p.model.add_subsystem('comp', om.ExplicitFuncComp(f))
p.setup(mode='fwd')
p.run_model()
assert_check_totals(p.check_totals(of=['comp.x', 'comp.y'], wrt=['comp.a', 'comp.b', 'comp.c'], method='cs', out_stream=None))
p.setup(mode='rev')
p.run_model()
assert_check_totals(p.check_totals(of=['comp.x', 'comp.y'], wrt=['comp.a', 'comp.b', 'comp.c'], method='cs', out_stream=None))
def test_nested_pars(self):
p = om.Problem()
par = p.model.add_subsystem('par', om.ParallelGroup(), promotes_inputs=['x'])
G1 = par.add_subsystem('G1', om.Group(), promotes_inputs=['x'])
G1p = G1.add_subsystem('G1p', om.ParallelGroup(), promotes_inputs=['x'])
G1p.add_subsystem('C1_1', om.ExecComp('y = 3.0*x'))
G1p.add_subsystem('C1_2', om.ExecComp('y = -3.0*x'))
G1p.promotes('C1_1', inputs=['x'], src_indices=[0])
G1p.promotes('C1_2', inputs=['x'], src_indices=[0])
G2 = par.add_subsystem('G2', om.Group(), promotes_inputs=['x'])
G2p = G2.add_subsystem('G2p', om.ParallelGroup(), promotes_inputs=['x'])
G2p.add_subsystem('C2_1', om.ExecComp('y = 5.0*x'))
G2p.add_subsystem('C2_2', om.ExecComp('y = -5.0*x'))
G2p.promotes('C2_1', inputs=['x'], src_indices=[1])
G2p.promotes('C2_2', inputs=['x'], src_indices=[1])
G3 = par.add_subsystem('G3', om.Group(), promotes_inputs=['x'], max_procs=1) # no nested parallel group here
G3s = G3.add_subsystem('G3s', om.Group(), promotes_inputs=['x'])
G3s.add_subsystem('C3_1', om.ExecComp('y = 7.0*x'))
G3s.add_subsystem('C3_2', om.ExecComp('y = -7.0*x'))
G3s.promotes('C3_1', inputs=['x'], src_indices=[2])
G3s.promotes('C3_2', inputs=['x'], src_indices=[2])
par.set_input_defaults('x', val=[.5, 1.5, 2.5])
p.setup()
p.run_model()
assert_check_totals(p.check_totals(of=['par.G1.G1p.C1_1.y', 'par.G1.G1p.C1_2.y',
'par.G2.G2p.C2_1.y', 'par.G2.G2p.C2_2.y',
'par.G3.G3s.C3_1.y', 'par.G3.G3s.C3_2.y'], wrt=['x']))
def test_user_compute_partials_func(self):
def J_func(x, y, z, J):
# the following sub-jacs are 4x4 based on the sizes of foo, bar, x, and y, but the partials
# were declared specifying rows and cols (in this case sub-jacs are diagonal), so we only
# store the nonzero values of the sub-jacs, resulting in an actual size of 4 rather than 4x4.
J['foo', 'x'] = -3*np.log(z)/(3*x+2*y)**2
J['foo', 'y'] = -2*np.log(z)/(3*x+2*y)**2
J['bar', 'x'] = 2.*np.ones(4)
J['bar', 'y'] = np.ones(4)
# z is a scalar so the true size of this sub-jac is 4x1
J['foo', 'z'] = 1/(z*(3*x+2*y))
def func(x=np.zeros(4), y=np.ones(4), z=3):
foo = np.log(z)/(3*x+2*y)
bar = 2.*x + y
return foo, bar
f = (omf.wrap(func)
.defaults(units='m')
.add_output('foo', units='1/m', shape=4)
.add_output('bar', shape=4)
.declare_partials(of='foo', wrt=('x', 'y'), rows=np.arange(4), cols=np.arange(4))
.declare_partials(of='foo', wrt='z')
.declare_partials(of='bar', wrt=('x', 'y'), rows=np.arange(4), cols=np.arange(4)))
p = om.Problem()
p.model.add_subsystem('comp', om.ExplicitFuncComp(f, compute_partials=J_func))
p.setup(force_alloc_complex=True)
p.run_model()
assert_check_totals(p.check_totals(of=['comp.foo', 'comp.bar'], wrt=['comp.x', 'comp.y', 'comp.z'], method='cs'))
def test_solve_nonlinear(self):
def apply_nl(a, b, c, x):
R_x = a * x**2 + b * x + c
return R_x
def solve_nl(a, b, c, x):
x = (-b + (b**2 - 4 * a * c)**0.5) / (2 * a)
return x
f = (omf.wrap(apply_nl).add_output(
'x', resid='R_x', val=0.0).declare_partials(of='*',
wrt='*',
method='cs'))
p = om.Problem()
p.model.add_subsystem('comp',
om.ImplicitFuncComp(f, solve_nonlinear=solve_nl))
# need this since comp is implicit and doesn't have a solve_linear
p.model.linear_solver = om.DirectSolver()
p.setup()
p.set_val('comp.a', 2.)
p.set_val('comp.b', -8.)
p.set_val('comp.c', 6.)
p.run_model()
assert_check_partials(p.check_partials(includes=['comp'],
out_stream=None),
atol=1e-5)
assert_check_totals(
p.check_totals(of=['comp.x'],
wrt=['comp.a', 'comp.b', 'comp.c'],
out_stream=None))
def test_apply_nonlinear_option(self):
def apply_nl(a, b, c, x, opt):
R_x = a * x**2 + b * x + c
if opt == 'foo':
R_x = -R_x
return R_x
f = (omf.wrap(apply_nl).add_output(
'x', resid='R_x', val=0.0).declare_option(
'opt', default='foo').declare_partials(of='*',
wrt='*',
method='cs'))
p = om.Problem()
p.model.add_subsystem('comp', om.ImplicitFuncComp(f))
# need this since comp is implicit and doesn't have a solve_linear
p.model.linear_solver = om.DirectSolver()
p.model.nonlinear_solver = om.NewtonSolver(solve_subsystems=False,
iprint=0)
p.setup()
p.set_val('comp.a', 2.)
p.set_val('comp.b', -8.)
p.set_val('comp.c', 6.)
p.run_model()
assert_check_partials(p.check_partials(includes=['comp'],
out_stream=None),
atol=1e-5)
assert_check_totals(
p.check_totals(of=['comp.x'],
wrt=['comp.a', 'comp.b', 'comp.c'],
out_stream=None))
def test_apply_nonlinear_linsys_coloring_cs(self):
prob = self.setup_apply_nonlinear_linsys_coloring('cs', 'fwd')
partials = prob.check_partials(includes=['comp'], out_stream=None)
assert_check_partials(partials, atol=1e-5)
assert_check_totals(prob.check_totals(of=['comp.x'],
wrt=['comp.A', 'comp.b'],
out_stream=None),
atol=3e-5,
rtol=3e-5)
def test_solve_lin_nl_linearize_reordered_args(self):
def apply_nl(x, a, b, c):
R_x = a * x**2 + b * x + c
return R_x
def solve_nl(x, a, b, c):
x = (-b + (b**2 - 4 * a * c)**0.5) / (2 * a)
return x
def linearize(x, a, b, c, partials):
partials['x', 'a'] = x**2
partials['x', 'b'] = x
partials['x', 'c'] = 1.0
partials['x', 'x'] = 2 * a * x + b
inv_jac = 1.0 / (2 * a * x + b)
return inv_jac
def solve_linear(d_x, mode, inv_jac):
if mode == 'fwd':
d_x = inv_jac * d_x
return d_x
elif mode == 'rev':
dR_x = inv_jac * d_x
return dR_x
f = (omf.wrap(apply_nl).add_output('x', resid='R_x',
val=0.0).declare_partials(of='*',
wrt='*'))
p = om.Problem()
p.model.add_subsystem(
'comp',
om.ImplicitFuncComp(f,
solve_linear=solve_linear,
linearize=linearize,
solve_nonlinear=solve_nl))
p.setup()
p.set_val('comp.a', 2.)
p.set_val('comp.b', -8.)
p.set_val('comp.c', 6.)
p.run_model()
assert_check_partials(p.check_partials(includes=['comp'],
out_stream=None),
atol=1e-5)
assert_check_totals(
p.check_totals(of=['comp.x'],
wrt=['comp.a', 'comp.b', 'comp.c'],
out_stream=None))
def test_coloring2(self):
# this test uses a very narrow matrix so the attempt at partial coloring will abort.
# There used to be a bug where the total derivatives would be incorrect when this
# happened.
mat = np.array([[0.14898778],
[0.19860233],
[0.81899035],
[0.78498818],
[0.68436335],
[0.93677595],
[0.33964473],
[0.82057559],
[0.62672187],
[0.52089597],
[0.28524249],
[0.62003238]])
inshapes = [1]
outshapes = [2, 3, 7]
outsizes = [np.prod(shp) for shp in outshapes]
def func(a):
ovec = mat.dot(a.flat)
x, y, z = ovec2outs(ovec, outsizes)
return x, y, z
f = (omf.wrap(func)
.add_inputs(a={'shape': inshapes[0]})
.add_outputs(x={'shape': outshapes[0]}, y={'shape': outshapes[1]}, z={'shape': outshapes[2]})
.declare_coloring(wrt='*', method='cs', show_summary=False)
)
p = om.Problem()
p.model.add_subsystem('comp', om.ExplicitFuncComp(f))
p.setup(mode='fwd')
p.run_model()
assert_check_totals(p.check_totals(of=['comp.x', 'comp.y', 'comp.z'], wrt=['comp.a'], out_stream=None))
p.setup(mode='rev')
p.run_model()
assert_check_totals(p.check_totals(of=['comp.x', 'comp.y', 'comp.z'], wrt=['comp.a'], out_stream=None))
def test_multipoint2(self):
p = om.Problem()
par = p.model.add_subsystem('par', om.ParallelGroup())
g1 = par.add_subsystem('g1', om.Group(), promotes_inputs=['x'])
g1.add_subsystem('C1', om.ExecComp('y = 3*x', shape=1))
g1.promotes('C1', inputs=['x'], src_indices=[0], src_shape=(3,))
g2 = par.add_subsystem('g2', om.Group(), promotes_inputs=['x'])
g2.add_subsystem('C2', om.ExecComp('y = 2*x', shape=1))
g2.promotes('C2', inputs=['x'], src_indices=[1], src_shape=(3,))
g3 = par.add_subsystem('g3', om.Group(), promotes_inputs=['x'])
g3.add_subsystem('C3', om.ExecComp('y = 5*x', shape=1))
g3.promotes('C3', inputs=['x'], src_indices=[2], src_shape=(3,))
p.model.set_input_defaults('par.x', val=[7., -5., 2.])
p.setup()
p.run_model()
assert_check_totals(p.check_totals(of=['par.g1.C1.y', 'par.g2.C2.y', 'par.g3.C3.y'], wrt=['par.x']))