def test_noncontiguous_idxs(self): # take even input indices in 0 rank and odd ones in 1 rank size = 11 p = Problem(model=Group()) top = p.model C1 = top.add_subsystem("C1", InOutArrayComp(size)) C2 = top.add_subsystem("C2", DistribNoncontiguousComp(size)) C3 = top.add_subsystem("C3", DistribGatherComp(size)) top.connect('C1.outvec', 'C2.invec') top.connect('C2.outvec', 'C3.invec') p.setup(check=False) # Conclude setup but don't run model. p.final_setup() C1._inputs['invec'] = np.array(range(size), float) p.run_model() if MPI: if self.comm.rank == 0: self.assertTrue(all(C2._outputs['outvec'] == np.array(list(take_nth(0, 2, range(size))), 'f')*4)) else: self.assertTrue(all(C2._outputs['outvec'] == np.array(list(take_nth(1, 2, range(size))), 'f')*4)) full_list = list(take_nth(0, 2, range(size))) + list(take_nth(1, 2, range(size))) self.assertTrue(all(C3._outputs['outvec'] == np.array(full_list, 'f')*4)) else: self.assertTrue(all(C2._outputs['outvec'] == C1._outputs['outvec']*2.)) self.assertTrue(all(C3._outputs['outvec'] == C2._outputs['outvec']))
def setup(self): comm = self.comm rank = comm.rank idxs = list(take_nth(rank, comm.size, range(self.arr_size))) self.add_input('invec', np.ones(len(idxs), float), src_indices=idxs) self.add_output('outvec', np.ones(len(idxs), float))
def setup(self): comm = self.comm rank = comm.rank arr_size = self.options['arr_size'] idxs = list(take_nth(rank, comm.size, range(arr_size))) self.add_input('invec', np.ones(len(idxs), float), distributed=True, src_indices=idxs) self.add_output('outvec', np.ones(len(idxs), float), distributed=True)