47 #include "Teko_InterlacedTpetra.hpp" 48 #include "Tpetra_Import.hpp" 56 namespace TpetraHelpers {
62 void buildSubMaps(GO numGlobals,
int numVars,
const Teuchos::Comm<int> & comm,std::vector<std::pair<
int,RCP<Tpetra::Map<LO,GO,NT> > > > & subMaps)
64 std::vector<int> vars;
67 for(
int i=0;i<numVars;i++) vars.push_back(1);
70 buildSubMaps(numGlobals,vars,comm,subMaps);
74 void buildSubMaps(
const Tpetra::Map<LO,GO,NT> & globalMap,
const std::vector<int> & vars,
const Teuchos::Comm<int> & comm,
75 std::vector<std::pair<
int,Teuchos::RCP<Tpetra::Map<LO,GO,NT> > > > & subMaps)
77 buildSubMaps(globalMap.getGlobalNumElements(),globalMap.getNodeNumElements(),globalMap.getMinGlobalIndex(),
82 void buildSubMaps(GO numGlobals,
const std::vector<int> & vars,
const Teuchos::Comm<int> & comm,std::vector<std::pair<
int,Teuchos::RCP<Tpetra::Map<LO,GO,NT> > > > & subMaps)
84 std::vector<int>::const_iterator varItr;
87 int numGlobalVars = 0;
88 for(varItr=vars.begin();varItr!=vars.end();++varItr)
89 numGlobalVars += *varItr;
92 TEUCHOS_ASSERT((numGlobals%numGlobalVars)==0);
94 Tpetra::Map<LO,GO,NT> sampleMap(numGlobals/numGlobalVars,0,rcpFromRef(comm));
96 buildSubMaps(numGlobals,numGlobalVars*sampleMap.getNodeNumElements(),numGlobalVars*sampleMap.getMinGlobalIndex(),vars,comm,subMaps);
100 void buildSubMaps(GO numGlobals,LO numMyElements,GO minMyGID,
const std::vector<int> & vars,
const Teuchos::Comm<int> & comm,
101 std::vector<std::pair<
int,Teuchos::RCP<Tpetra::Map<LO,GO,NT> > > > & subMaps)
103 std::vector<int>::const_iterator varItr;
106 int numGlobalVars = 0;
107 for(varItr=vars.begin();varItr!=vars.end();++varItr)
108 numGlobalVars += *varItr;
111 TEUCHOS_ASSERT((numGlobals%numGlobalVars)==0);
112 TEUCHOS_ASSERT((numMyElements%numGlobalVars)==0);
113 TEUCHOS_ASSERT((minMyGID%numGlobalVars)==0);
115 LO numBlocks = numMyElements/numGlobalVars;
116 GO minBlockID = minMyGID/numGlobalVars;
122 for(varItr=vars.begin();varItr!=vars.end();++varItr) {
123 LO numLocalVars = *varItr;
124 GO numAllElmts = numLocalVars*numGlobals/numGlobalVars;
125 LO numMyElmts = numLocalVars * numBlocks;
128 std::vector<GO> subGlobals;
129 std::vector<GO> contigGlobals;
133 for(LO blockNum=0;blockNum<numBlocks;blockNum++) {
136 for(LO local=0;local<numLocalVars;++local) {
140 subGlobals.push_back((minBlockID+blockNum)*numGlobalVars+blockOffset+local);
143 contigGlobals.push_back(numLocalVars*minBlockID+count);
149 assert((
size_t) numMyElmts==subGlobals.size());
152 RCP<Tpetra::Map<LO,GO,NT> > subMap = rcp(
new Tpetra::Map<LO,GO,NT>(numAllElmts,Teuchos::ArrayView<GO>(subGlobals),0,rcpFromRef(comm)));
153 RCP<Tpetra::Map<LO,GO,NT> > contigMap = rcp(
new Tpetra::Map<LO,GO,NT>(numAllElmts,Teuchos::ArrayView<GO>(contigGlobals),0,rcpFromRef(comm)));
155 Teuchos::set_extra_data(contigMap,
"contigMap",Teuchos::inOutArg(subMap));
156 subMaps.push_back(std::make_pair(numLocalVars,subMap));
159 blockOffset += numLocalVars;
163 void buildExportImport(
const Tpetra::Map<LO,GO,NT> & baseMap,
const std::vector<std::pair<
int,RCP<Tpetra::Map<LO,GO,NT> > > > & subMaps,
164 std::vector<RCP<Tpetra::Export<LO,GO,NT> > > & subExport,
165 std::vector<RCP<Tpetra::Import<LO,GO,NT> > > & subImport)
167 std::vector<std::pair<int,RCP<Tpetra::Map<LO,GO,NT> > > >::const_iterator mapItr;
170 for(mapItr=subMaps.begin();mapItr!=subMaps.end();++mapItr) {
172 const Tpetra::Map<LO,GO,NT> & map = *(mapItr->second);
175 subImport.push_back(rcp(
new Tpetra::Import<LO,GO,NT>(rcpFromRef(baseMap),rcpFromRef(map))));
176 subExport.push_back(rcp(
new Tpetra::Export<LO,GO,NT>(rcpFromRef(map),rcpFromRef(baseMap))));
180 void buildSubVectors(
const std::vector<std::pair<
int,RCP<Tpetra::Map<LO,GO,NT> > > > & subMaps,std::vector<RCP<Tpetra::MultiVector<ST,LO,GO,NT> > > & subVectors,
int count)
182 std::vector<std::pair<int,RCP<Tpetra::Map<LO,GO,NT> > > >::const_iterator mapItr;
185 for(mapItr=subMaps.begin();mapItr!=subMaps.end();++mapItr) {
187 const Tpetra::Map<LO,GO,NT> & map = *(Teuchos::get_extra_data<RCP<Tpetra::Map<LO,GO,NT> > >(mapItr->second,
"contigMap"));
190 RCP<Tpetra::MultiVector<ST,LO,GO,NT> > mv = rcp(
new Tpetra::MultiVector<ST,LO,GO,NT>(rcpFromRef(map),count));
191 Teuchos::set_extra_data(mapItr->second,
"globalMap",Teuchos::inOutArg(mv));
192 subVectors.push_back(mv);
196 void associateSubVectors(
const std::vector<std::pair<
int,RCP<Tpetra::Map<LO,GO,NT> > > > & subMaps,std::vector<RCP<
const Tpetra::MultiVector<ST,LO,GO,NT> > > & subVectors)
198 std::vector<std::pair<int,RCP<Tpetra::Map<LO,GO,NT> > > >::const_iterator mapItr;
199 std::vector<RCP<const Tpetra::MultiVector<ST,LO,GO,NT> > >::iterator vecItr;
201 TEUCHOS_ASSERT(subMaps.size()==subVectors.size());
204 for(mapItr=subMaps.begin(),vecItr=subVectors.begin();mapItr!=subMaps.end();++mapItr,++vecItr)
205 Teuchos::set_extra_data(mapItr->second,
"globalMap",Teuchos::inOutArg(*vecItr),Teuchos::POST_DESTROY,
false);
209 RCP<Tpetra::CrsMatrix<ST,LO,GO,NT> > buildSubBlock(
int i,
int j,
const RCP<
const Tpetra::CrsMatrix<ST,LO,GO,NT> >& A,
const std::vector<std::pair<
int,RCP<Tpetra::Map<LO,GO,NT> > > > & subMaps)
212 int numVarFamily = subMaps.size();
214 TEUCHOS_ASSERT(i>=0 && i<numVarFamily);
215 TEUCHOS_ASSERT(j>=0 && j<numVarFamily);
217 const Tpetra::Map<LO,GO,NT> & gRowMap = *subMaps[i].second;
218 const Tpetra::Map<LO,GO,NT> & rowMap = *Teuchos::get_extra_data<RCP<Tpetra::Map<LO,GO,NT> > >(subMaps[i].second,
"contigMap");
219 const Tpetra::Map<LO,GO,NT> & colMap = *Teuchos::get_extra_data<RCP<Tpetra::Map<LO,GO,NT> > >(subMaps[j].second,
"contigMap");
220 int colFamilyCnt = subMaps[j].first;
224 GO numGlobalVars = 0;
225 GO rowBlockOffset = 0;
226 GO colBlockOffset = 0;
227 for(
int k=0;k<numVarFamily;k++) {
228 numGlobalVars += subMaps[k].first;
231 if(k<i) rowBlockOffset += subMaps[k].first;
232 if(k<j) colBlockOffset += subMaps[k].first;
236 Tpetra::Import<LO,GO,NT>
import(A->getRowMap(),rcpFromRef(gRowMap));
237 Tpetra::CrsMatrix<ST,LO,GO,NT> localA(rcpFromRef(gRowMap),0);
238 localA.doImport(*A,
import,Tpetra::INSERT);
240 RCP<Tpetra::CrsMatrix<ST,LO,GO,NT> > mat = Tpetra::createCrsMatrix<ST,LO,GO,NT>(rcpFromRef(rowMap),0);
243 LO numMyRows = rowMap.getNodeNumElements();
244 LO maxNumEntries = A->getGlobalMaxNumRowEntries();
247 std::vector<GO> indices(maxNumEntries);
248 std::vector<ST> values(maxNumEntries);
251 std::vector<GO> colIndices(maxNumEntries);
252 std::vector<ST> colValues(maxNumEntries);
256 for(LO localRow=0;localRow<numMyRows;localRow++) {
257 size_t numEntries = -1;
258 GO globalRow = gRowMap.getGlobalElement(localRow);
259 GO contigRow = rowMap.getGlobalElement(localRow);
261 TEUCHOS_ASSERT(globalRow>=0);
262 TEUCHOS_ASSERT(contigRow>=0);
265 localA.getGlobalRowCopy(globalRow, Teuchos::ArrayView<GO>(indices), Teuchos::ArrayView<ST>(values), numEntries);
267 for(
size_t localCol=0;localCol<numEntries;localCol++) {
268 GO globalCol = indices[localCol];
271 int block = globalCol / numGlobalVars;
273 bool inFamily =
true;
276 inFamily &= (block*numGlobalVars+colBlockOffset <= globalCol);
277 inFamily &= ((block*numGlobalVars+colBlockOffset+colFamilyCnt) > globalCol);
281 GO familyOffset = globalCol-(block*numGlobalVars+colBlockOffset);
283 colIndices[numOwnedCols] = block*colFamilyCnt + familyOffset;
284 colValues[numOwnedCols] = values[localCol];
291 colIndices.resize(numOwnedCols);
292 colValues.resize(numOwnedCols);
293 mat->insertGlobalValues(contigRow,Teuchos::ArrayView<GO>(colIndices),Teuchos::ArrayView<ST>(colValues));
294 colIndices.resize(maxNumEntries);
295 colValues.resize(maxNumEntries);
299 mat->fillComplete(rcpFromRef(colMap),rcpFromRef(rowMap));
305 void rebuildSubBlock(
int i,
int j,
const RCP<
const Tpetra::CrsMatrix<ST,LO,GO,NT> > & A,
const std::vector<std::pair<
int,RCP<Tpetra::Map<LO,GO,NT> > > > & subMaps,Tpetra::CrsMatrix<ST,LO,GO,NT> & mat)
308 int numVarFamily = subMaps.size();
310 TEUCHOS_ASSERT(i>=0 && i<numVarFamily);
311 TEUCHOS_ASSERT(j>=0 && j<numVarFamily);
312 TEUCHOS_ASSERT(mat.isFillComplete());
314 const Tpetra::Map<LO,GO,NT> & gRowMap = *subMaps[i].second;
315 const Tpetra::Map<LO,GO,NT> & rowMap = *Teuchos::get_extra_data<RCP<Tpetra::Map<LO,GO,NT> > >(subMaps[i].second,
"contigMap");
316 const Tpetra::Map<LO,GO,NT> & colMap = *Teuchos::get_extra_data<RCP<Tpetra::Map<LO,GO,NT> > >(subMaps[j].second,
"contigMap");
317 int colFamilyCnt = subMaps[j].first;
321 GO numGlobalVars = 0;
322 GO rowBlockOffset = 0;
323 GO colBlockOffset = 0;
324 for(
int k=0;k<numVarFamily;k++) {
325 numGlobalVars += subMaps[k].first;
328 if(k<i) rowBlockOffset += subMaps[k].first;
329 if(k<j) colBlockOffset += subMaps[k].first;
333 Tpetra::Import<LO,GO,NT>
import(A->getRowMap(),rcpFromRef(gRowMap));
334 Tpetra::CrsMatrix<ST,LO,GO,NT> localA(rcpFromRef(gRowMap),0);
335 localA.doImport(*A,
import,Tpetra::INSERT);
339 mat.setAllToScalar(0.0);
342 LO numMyRows = rowMap.getNodeNumElements();
343 GO maxNumEntries = A->getGlobalMaxNumRowEntries();
346 std::vector<GO> indices(maxNumEntries);
347 std::vector<ST> values(maxNumEntries);
350 std::vector<GO> colIndices(maxNumEntries);
351 std::vector<ST> colValues(maxNumEntries);
355 for(LO localRow=0;localRow<numMyRows;localRow++) {
356 size_t numEntries = -1;
357 GO globalRow = gRowMap.getGlobalElement(localRow);
358 GO contigRow = rowMap.getGlobalElement(localRow);
360 TEUCHOS_ASSERT(globalRow>=0);
361 TEUCHOS_ASSERT(contigRow>=0);
364 localA.getGlobalRowCopy(globalRow, Teuchos::ArrayView<GO>(indices), Teuchos::ArrayView<ST>(values), numEntries);
367 for(
size_t localCol=0;localCol<numEntries;localCol++) {
368 GO globalCol = indices[localCol];
371 int block = globalCol / numGlobalVars;
373 bool inFamily =
true;
376 inFamily &= (block*numGlobalVars+colBlockOffset <= globalCol);
377 inFamily &= ((block*numGlobalVars+colBlockOffset+colFamilyCnt) > globalCol);
381 GO familyOffset = globalCol-(block*numGlobalVars+colBlockOffset);
383 colIndices[numOwnedCols] = block*colFamilyCnt + familyOffset;
384 colValues[numOwnedCols] = values[localCol];
391 colIndices.resize(numOwnedCols);
392 colValues.resize(numOwnedCols);
393 mat.sumIntoGlobalValues(contigRow,Teuchos::ArrayView<GO>(colIndices),Teuchos::ArrayView<ST>(colValues));
394 colIndices.resize(maxNumEntries);
395 colValues.resize(maxNumEntries);
397 mat.fillComplete(rcpFromRef(colMap),rcpFromRef(rowMap));
402 void many2one(Tpetra::MultiVector<ST,LO,GO,NT> & one,
const std::vector<RCP<
const Tpetra::MultiVector<ST,LO,GO,NT> > > & many,
403 const std::vector<RCP<Tpetra::Export<LO,GO,NT> > > & subExport)
406 std::vector<RCP<const Tpetra::MultiVector<ST,LO,GO,NT> > >::const_iterator vecItr;
407 std::vector<RCP<Tpetra::Export<LO,GO,NT> > >::const_iterator expItr;
410 for(vecItr=many.begin(),expItr=subExport.begin();
411 vecItr!=many.end();++vecItr,++expItr) {
414 RCP<const Tpetra::MultiVector<ST,LO,GO,NT> > srcVec = *vecItr;
417 const Tpetra::Map<LO,GO,NT> & globalMap = *(Teuchos::get_extra_data<RCP<Tpetra::Map<LO,GO,NT> > >(srcVec,
"globalMap"));
420 GO lda = srcVec->getStride();
421 GO srcSize = srcVec->getGlobalLength()*srcVec->getNumVectors();
422 std::vector<ST> srcArray(srcSize);
423 Teuchos::ArrayView<ST> srcVals(srcArray);
424 srcVec->get1dCopy(srcVals,lda);
425 Tpetra::MultiVector<ST,LO,GO,NT> exportVector(rcpFromRef(globalMap),srcVals,lda,srcVec->getNumVectors());
428 one.doExport(exportVector,**expItr,Tpetra::INSERT);
433 void one2many(std::vector<RCP<Tpetra::MultiVector<ST,LO,GO,NT> > > & many,
const Tpetra::MultiVector<ST,LO,GO,NT> & single,
434 const std::vector<RCP<Tpetra::Import<LO,GO,NT> > > & subImport)
437 std::vector<RCP<Tpetra::MultiVector<ST,LO,GO,NT> > >::const_iterator vecItr;
438 std::vector<RCP<Tpetra::Import<LO,GO,NT> > >::const_iterator impItr;
441 for(vecItr=many.begin(),impItr=subImport.begin();
442 vecItr!=many.end();++vecItr,++impItr) {
444 RCP<Tpetra::MultiVector<ST,LO,GO,NT> > destVec = *vecItr;
447 const Tpetra::Map<LO,GO,NT> & globalMap = *(Teuchos::get_extra_data<RCP<Tpetra::Map<LO,GO,NT> > >(destVec,
"globalMap"));
450 GO destLDA = destVec->getStride();
451 GO destSize = destVec->getGlobalLength()*destVec->getNumVectors();
452 std::vector<ST> destArray(destSize);
453 Teuchos::ArrayView<ST> destVals(destArray);
454 destVec->get1dCopy(destVals,destLDA);
455 Tpetra::MultiVector<ST,LO,GO,NT> importVector(rcpFromRef(globalMap),destVals,destLDA,destVec->getNumVectors());
458 importVector.doImport(single,**impItr,Tpetra::INSERT);
460 Tpetra::Import<LO,GO,NT> importer(destVec->getMap(),destVec->getMap());
461 importVector.replaceMap(destVec->getMap());
462 destVec->doImport(importVector,importer,Tpetra::INSERT);