def test_user_partials_sparse(self):
def func3(x1=np.ones(5), x2=np.ones(5)):
y1=2.0*x1+1.
y2=3.0*x2-1.
return y1, y2
f = (omf.wrap(func3)
.defaults(shape=5)
.declare_partials(of=['y1'], wrt=['x1'], method='cs',
rows=np.arange(5), cols=np.arange(5))
.declare_partials(of=['y2'], wrt=['x2'], method='cs',
rows=np.arange(5), cols=np.arange(5)))
p = om.Problem()
p.model.add_subsystem('comp', ExplicitFuncCompCountRuns(f))
p.setup(mode='fwd')
p.run_driver()
assert_check_partials(p.check_partials(includes=['comp'], out_stream=None))
coloring = compute_total_coloring(p, mode='fwd', of=['comp.y1', 'comp.y2'], wrt=['comp.x1', 'comp.x2'], use_abs_names=True)
self.assertEqual(coloring.total_solves(), 1) # verify that sparsity had an effect
def color(self, mode='auto', fname=None):
self.coloring = compute_total_coloring(None, mode, setup=False,
run_model=False, bool_jac=self.J,
fname=fname)
return self.coloring
def test_totals_over_implicit_comp(self, method, isplit, osplit):
prob = Problem(coloring_dir=self.tempdir)
model = prob.model = Group()
sparsity = setup_sparsity(_BIGMASK)
indeps, conns = setup_indeps(isplit, _BIGMASK.shape[1], 'indeps',
'comp')
model.nonlinear_solver = NonlinearBlockGS()
model.add_subsystem('indeps', indeps)
comp = model.add_subsystem(
'comp',
SparseCompImplicit(sparsity, method, isplit=isplit, osplit=osplit))
for conn in conns:
model.connect(*conn)
for i in range(isplit):
model.add_design_var('indeps.x%d' % i)
for i in range(osplit):
model.comp.add_constraint('y%d' % i)
model.approx_totals(method=method)
prob.setup(check=False, mode='fwd')
prob.set_solver_print(level=0)
prob.run_model()
model._save_coloring(compute_total_coloring(prob))
prob = Problem(coloring_dir=self.tempdir)
model = prob.model = Group()
indeps, conns = setup_indeps(isplit, _BIGMASK.shape[1], 'indeps',
'comp')
model.nonlinear_solver = NonlinearBlockGS()
model.add_subsystem('indeps', indeps)
comp = model.add_subsystem(
'comp',
SparseCompImplicit(sparsity, method, isplit=isplit, osplit=osplit))
for conn in conns:
model.connect(*conn)
for i in range(isplit):
model.add_design_var('indeps.x%d' % i)
for i in range(osplit):
model.comp.add_constraint('y%d' % i)
model.approx_totals(method=method)
model.declare_coloring(wrt='*', method=method)
model.use_fixed_coloring()
prob.setup(check=False, mode='fwd')
prob.set_solver_print(level=0)
prob.run_model()
start_nruns = comp._nruns
derivs = prob.driver._compute_totals() # colored
nruns = comp._nruns - start_nruns
# NLBGS ends up doing a single iteration after initialization, resulting in 2
# runs per NL solve, so we multiplly the number of colored solvers by 2
self.assertEqual(nruns, 10 * 2)
_check_total_matrix(model, derivs, sparsity, method)
def test_totals_of_indices(self, method):
prob = Problem(coloring_dir=self.tempdir)
model = prob.model = Group()
mask = np.array([[1, 0, 0, 1, 1], [0, 1, 0, 1, 1], [0, 1, 0, 1, 1],
[1, 0, 0, 0, 0], [0, 1, 1, 0, 0]])
isplit = 2
sparsity = setup_sparsity(mask)
indeps, conns = setup_indeps(isplit, mask.shape[1], 'indeps', 'comp')
model.add_subsystem('indeps', indeps)
comp = model.add_subsystem(
'comp',
SparseCompExplicit(sparsity, method, isplit=isplit, osplit=2))
for conn in conns:
model.connect(*conn)
model.comp.add_constraint('y0', indices=[0, 2])
model.comp.add_constraint('y1')
model.add_design_var('indeps.x0')
model.add_design_var('indeps.x1')
model.approx_totals(method=method)
prob.setup(check=False, mode='fwd')
prob.set_solver_print(level=0)
prob.run_model()
model._save_coloring(compute_total_coloring(prob))
prob = Problem(coloring_dir=self.tempdir)
model = prob.model = Group()
indeps, conns = setup_indeps(isplit, mask.shape[1], 'indeps', 'comp')
model.add_subsystem('indeps', indeps)
comp = model.add_subsystem(
'comp',
SparseCompExplicit(sparsity, method, isplit=isplit, osplit=2))
model.connect('indeps.x0', 'comp.x0')
model.connect('indeps.x1', 'comp.x1')
model.comp.add_constraint('y0', indices=[0, 2])
model.comp.add_constraint('y1')
model.add_design_var('indeps.x0')
model.add_design_var('indeps.x1')
model.approx_totals(method=method)
model.declare_coloring(wrt='*', method=method)
model.use_fixed_coloring()
prob.setup(check=False, mode='fwd')
prob.set_solver_print(level=0)
prob.run_model()
start_nruns = comp._nruns
derivs = prob.driver._compute_totals() # colored
nruns = comp._nruns - start_nruns
self.assertEqual(nruns, 3)
rows = [0, 2, 3, 4]
_check_total_matrix(model, derivs, sparsity[rows, :], method)
def test_simple_totals_static(self, method, isplit, osplit):
prob = Problem(coloring_dir=self.tempdir)
model = prob.model = Group()
sparsity = setup_sparsity(_BIGMASK)
indeps, conns = setup_indeps(isplit, _BIGMASK.shape[1], 'indeps',
'comp')
model.add_subsystem('indeps', indeps)
comp = model.add_subsystem(
'comp',
SparseCompExplicit(sparsity, method, isplit=isplit, osplit=osplit))
for conn in conns:
model.connect(*conn)
for i in range(isplit):
model.add_design_var('indeps.x%d' % i)
for i in range(osplit):
model.comp.add_constraint('y%d' % i)
model.approx_totals(method=method)
prob.setup(check=False, mode='fwd')
prob.set_solver_print(level=0)
prob.run_model()
# try just manually computing the coloring here instead of using declare_coloring
coloring = compute_total_coloring(prob)
model._save_coloring(coloring)
# new Problem, loading the coloring we just computed
prob = Problem(coloring_dir=self.tempdir)
model = prob.model = Group()
indeps, conns = setup_indeps(isplit, _BIGMASK.shape[1], 'indeps',
'comp')
model.add_subsystem('indeps', indeps)
comp = model.add_subsystem(
'comp',
SparseCompExplicit(sparsity, method, isplit=isplit, osplit=osplit))
for conn in conns:
model.connect(*conn)
for i in range(isplit):
model.add_design_var('indeps.x%d' % i)
for i in range(osplit):
model.comp.add_constraint('y%d' % i)
model.approx_totals(method=method)
model.declare_coloring(wrt='*', method=method)
model.use_fixed_coloring()
prob.setup(check=False, mode='fwd')
prob.set_solver_print(level=0)
prob.run_model()
start_nruns = comp._nruns
derivs = prob.driver._compute_totals()
nruns = comp._nruns - start_nruns
self.assertEqual(nruns, 10)
_check_total_matrix(model, derivs, sparsity, method)