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)