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;
126 LO numMyElmts = numLocalVars * numBlocks;
130 std::vector<GO> subGlobals;
131 std::vector<GO> contigGlobals;
135 for(LO blockNum=0;blockNum<numBlocks;blockNum++) {
138 for(LO local=0;local<numLocalVars;++local) {
142 subGlobals.push_back((minBlockID+blockNum)*numGlobalVars+blockOffset+local);
145 contigGlobals.push_back(numLocalVars*minBlockID+count);
151 assert((
size_t) numMyElmts==subGlobals.size());
154 RCP<Tpetra::Map<LO,GO,NT> > subMap = rcp(
new Tpetra::Map<LO,GO,NT>(numAllElmts,Teuchos::ArrayView<GO>(subGlobals),0,rcpFromRef(comm)));
155 RCP<Tpetra::Map<LO,GO,NT> > contigMap = rcp(
new Tpetra::Map<LO,GO,NT>(numAllElmts,Teuchos::ArrayView<GO>(contigGlobals),0,rcpFromRef(comm)));
157 Teuchos::set_extra_data(contigMap,
"contigMap",Teuchos::inOutArg(subMap));
158 subMaps.push_back(std::make_pair(numLocalVars,subMap));
161 blockOffset += numLocalVars;
165 void buildExportImport(
const Tpetra::Map<LO,GO,NT> & baseMap,
const std::vector<std::pair<
int,RCP<Tpetra::Map<LO,GO,NT> > > > & subMaps,
166 std::vector<RCP<Tpetra::Export<LO,GO,NT> > > & subExport,
167 std::vector<RCP<Tpetra::Import<LO,GO,NT> > > & subImport)
169 std::vector<std::pair<int,RCP<Tpetra::Map<LO,GO,NT> > > >::const_iterator mapItr;
172 for(mapItr=subMaps.begin();mapItr!=subMaps.end();++mapItr) {
174 const Tpetra::Map<LO,GO,NT> & map = *(mapItr->second);
177 subImport.push_back(rcp(
new Tpetra::Import<LO,GO,NT>(rcpFromRef(baseMap),rcpFromRef(map))));
178 subExport.push_back(rcp(
new Tpetra::Export<LO,GO,NT>(rcpFromRef(map),rcpFromRef(baseMap))));
182 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)
184 std::vector<std::pair<int,RCP<Tpetra::Map<LO,GO,NT> > > >::const_iterator mapItr;
187 for(mapItr=subMaps.begin();mapItr!=subMaps.end();++mapItr) {
189 const Tpetra::Map<LO,GO,NT> & map = *(Teuchos::get_extra_data<RCP<Tpetra::Map<LO,GO,NT> > >(mapItr->second,
"contigMap"));
192 RCP<Tpetra::MultiVector<ST,LO,GO,NT> > mv = rcp(
new Tpetra::MultiVector<ST,LO,GO,NT>(rcpFromRef(map),count));
193 Teuchos::set_extra_data(mapItr->second,
"globalMap",Teuchos::inOutArg(mv));
194 subVectors.push_back(mv);
198 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)
200 std::vector<std::pair<int,RCP<Tpetra::Map<LO,GO,NT> > > >::const_iterator mapItr;
201 std::vector<RCP<const Tpetra::MultiVector<ST,LO,GO,NT> > >::iterator vecItr;
203 TEUCHOS_ASSERT(subMaps.size()==subVectors.size());
206 for(mapItr=subMaps.begin(),vecItr=subVectors.begin();mapItr!=subMaps.end();++mapItr,++vecItr)
207 Teuchos::set_extra_data(mapItr->second,
"globalMap",Teuchos::inOutArg(*vecItr),Teuchos::POST_DESTROY,
false);
211 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)
214 int numVarFamily = subMaps.size();
216 TEUCHOS_ASSERT(i>=0 && i<numVarFamily);
217 TEUCHOS_ASSERT(j>=0 && j<numVarFamily);
219 const Tpetra::Map<LO,GO,NT> & gRowMap = *subMaps[i].second;
220 const Tpetra::Map<LO,GO,NT> & rowMap = *Teuchos::get_extra_data<RCP<Tpetra::Map<LO,GO,NT> > >(subMaps[i].second,
"contigMap");
221 const Tpetra::Map<LO,GO,NT> & colMap = *Teuchos::get_extra_data<RCP<Tpetra::Map<LO,GO,NT> > >(subMaps[j].second,
"contigMap");
222 int colFamilyCnt = subMaps[j].first;
226 GO numGlobalVars = 0;
227 GO rowBlockOffset = 0;
228 GO colBlockOffset = 0;
229 for(
int k=0;k<numVarFamily;k++) {
230 numGlobalVars += subMaps[k].first;
233 if(k<i) rowBlockOffset += subMaps[k].first;
234 if(k<j) colBlockOffset += subMaps[k].first;
238 Tpetra::Import<LO,GO,NT>
import(A->getRowMap(),rcpFromRef(gRowMap));
239 Tpetra::CrsMatrix<ST,LO,GO,NT> localA(rcpFromRef(gRowMap),0);
240 localA.doImport(*A,
import,Tpetra::INSERT);
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<size_t> numEntriesPerRow(numMyRows,0);
254 for(LO localRow=0;localRow<numMyRows;localRow++) {
255 size_t numEntries = -1;
256 GO globalRow = gRowMap.getGlobalElement(localRow);
257 GO contigRow = rowMap.getGlobalElement(localRow);
259 TEUCHOS_ASSERT(globalRow>=0);
260 TEUCHOS_ASSERT(contigRow>=0);
263 localA.getGlobalRowCopy(globalRow, Teuchos::ArrayView<GO>(indices), Teuchos::ArrayView<ST>(values), numEntries);
265 for(
size_t localCol=0;localCol<numEntries;localCol++) {
266 GO globalCol = indices[localCol];
269 int block = globalCol / numGlobalVars;
271 bool inFamily =
true;
274 inFamily &= (block*numGlobalVars+colBlockOffset <= globalCol);
275 inFamily &= ((block*numGlobalVars+colBlockOffset+colFamilyCnt) > globalCol);
282 numEntriesPerRow[localRow] += numOwnedCols;
285 RCP<Tpetra::CrsMatrix<ST,LO,GO,NT> > mat =
286 rcp(
new Tpetra::CrsMatrix<ST,LO,GO,NT>(rcpFromRef(rowMap), Teuchos::ArrayView<const size_t>(numEntriesPerRow)));
289 std::vector<GO> colIndices(maxNumEntries);
290 std::vector<ST> colValues(maxNumEntries);
294 for(LO localRow=0;localRow<numMyRows;localRow++) {
295 size_t numEntries = -1;
296 GO globalRow = gRowMap.getGlobalElement(localRow);
297 GO contigRow = rowMap.getGlobalElement(localRow);
299 TEUCHOS_ASSERT(globalRow>=0);
300 TEUCHOS_ASSERT(contigRow>=0);
303 localA.getGlobalRowCopy(globalRow, Teuchos::ArrayView<GO>(indices), Teuchos::ArrayView<ST>(values), numEntries);
305 for(
size_t localCol=0;localCol<numEntries;localCol++) {
306 GO globalCol = indices[localCol];
309 int block = globalCol / numGlobalVars;
311 bool inFamily =
true;
314 inFamily &= (block*numGlobalVars+colBlockOffset <= globalCol);
315 inFamily &= ((block*numGlobalVars+colBlockOffset+colFamilyCnt) > globalCol);
319 GO familyOffset = globalCol-(block*numGlobalVars+colBlockOffset);
321 colIndices[numOwnedCols] = block*colFamilyCnt + familyOffset;
322 colValues[numOwnedCols] = values[localCol];
329 colIndices.resize(numOwnedCols);
330 colValues.resize(numOwnedCols);
331 mat->insertGlobalValues(contigRow,Teuchos::ArrayView<GO>(colIndices),Teuchos::ArrayView<ST>(colValues));
332 colIndices.resize(maxNumEntries);
333 colValues.resize(maxNumEntries);
337 mat->fillComplete(rcpFromRef(colMap),rcpFromRef(rowMap));
343 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)
346 int numVarFamily = subMaps.size();
348 TEUCHOS_ASSERT(i>=0 && i<numVarFamily);
349 TEUCHOS_ASSERT(j>=0 && j<numVarFamily);
350 TEUCHOS_ASSERT(mat.isFillComplete());
352 const Tpetra::Map<LO,GO,NT> & gRowMap = *subMaps[i].second;
353 const Tpetra::Map<LO,GO,NT> & rowMap = *Teuchos::get_extra_data<RCP<Tpetra::Map<LO,GO,NT> > >(subMaps[i].second,
"contigMap");
354 const Tpetra::Map<LO,GO,NT> & colMap = *Teuchos::get_extra_data<RCP<Tpetra::Map<LO,GO,NT> > >(subMaps[j].second,
"contigMap");
355 int colFamilyCnt = subMaps[j].first;
359 GO numGlobalVars = 0;
360 GO rowBlockOffset = 0;
361 GO colBlockOffset = 0;
362 for(
int k=0;k<numVarFamily;k++) {
363 numGlobalVars += subMaps[k].first;
366 if(k<i) rowBlockOffset += subMaps[k].first;
367 if(k<j) colBlockOffset += subMaps[k].first;
371 Tpetra::Import<LO,GO,NT>
import(A->getRowMap(),rcpFromRef(gRowMap));
372 Tpetra::CrsMatrix<ST,LO,GO,NT> localA(rcpFromRef(gRowMap),0);
373 localA.doImport(*A,
import,Tpetra::INSERT);
377 mat.setAllToScalar(0.0);
380 LO numMyRows = rowMap.getNodeNumElements();
381 GO maxNumEntries = A->getGlobalMaxNumRowEntries();
384 std::vector<GO> indices(maxNumEntries);
385 std::vector<ST> values(maxNumEntries);
388 std::vector<GO> colIndices(maxNumEntries);
389 std::vector<ST> colValues(maxNumEntries);
393 for(LO localRow=0;localRow<numMyRows;localRow++) {
394 size_t numEntries = -1;
395 GO globalRow = gRowMap.getGlobalElement(localRow);
396 GO contigRow = rowMap.getGlobalElement(localRow);
398 TEUCHOS_ASSERT(globalRow>=0);
399 TEUCHOS_ASSERT(contigRow>=0);
402 localA.getGlobalRowCopy(globalRow, Teuchos::ArrayView<GO>(indices), Teuchos::ArrayView<ST>(values), numEntries);
405 for(
size_t localCol=0;localCol<numEntries;localCol++) {
406 GO globalCol = indices[localCol];
409 int block = globalCol / numGlobalVars;
411 bool inFamily =
true;
414 inFamily &= (block*numGlobalVars+colBlockOffset <= globalCol);
415 inFamily &= ((block*numGlobalVars+colBlockOffset+colFamilyCnt) > globalCol);
419 GO familyOffset = globalCol-(block*numGlobalVars+colBlockOffset);
421 colIndices[numOwnedCols] = block*colFamilyCnt + familyOffset;
422 colValues[numOwnedCols] = values[localCol];
429 colIndices.resize(numOwnedCols);
430 colValues.resize(numOwnedCols);
431 mat.sumIntoGlobalValues(contigRow,Teuchos::ArrayView<GO>(colIndices),Teuchos::ArrayView<ST>(colValues));
432 colIndices.resize(maxNumEntries);
433 colValues.resize(maxNumEntries);
435 mat.fillComplete(rcpFromRef(colMap),rcpFromRef(rowMap));
440 void many2one(Tpetra::MultiVector<ST,LO,GO,NT> & one,
const std::vector<RCP<
const Tpetra::MultiVector<ST,LO,GO,NT> > > & many,
441 const std::vector<RCP<Tpetra::Export<LO,GO,NT> > > & subExport)
444 std::vector<RCP<const Tpetra::MultiVector<ST,LO,GO,NT> > >::const_iterator vecItr;
445 std::vector<RCP<Tpetra::Export<LO,GO,NT> > >::const_iterator expItr;
448 for(vecItr=many.begin(),expItr=subExport.begin();
449 vecItr!=many.end();++vecItr,++expItr) {
452 RCP<const Tpetra::MultiVector<ST,LO,GO,NT> > srcVec = *vecItr;
455 const Tpetra::Map<LO,GO,NT> & globalMap = *(Teuchos::get_extra_data<RCP<Tpetra::Map<LO,GO,NT> > >(srcVec,
"globalMap"));
458 GO lda = srcVec->getStride();
459 GO srcSize = srcVec->getGlobalLength()*srcVec->getNumVectors();
460 std::vector<ST> srcArray(srcSize);
461 Teuchos::ArrayView<ST> srcVals(srcArray);
462 srcVec->get1dCopy(srcVals,lda);
463 Tpetra::MultiVector<ST,LO,GO,NT> exportVector(rcpFromRef(globalMap),srcVals,lda,srcVec->getNumVectors());
466 one.doExport(exportVector,**expItr,Tpetra::INSERT);
471 void one2many(std::vector<RCP<Tpetra::MultiVector<ST,LO,GO,NT> > > & many,
const Tpetra::MultiVector<ST,LO,GO,NT> & single,
472 const std::vector<RCP<Tpetra::Import<LO,GO,NT> > > & subImport)
475 std::vector<RCP<Tpetra::MultiVector<ST,LO,GO,NT> > >::const_iterator vecItr;
476 std::vector<RCP<Tpetra::Import<LO,GO,NT> > >::const_iterator impItr;
479 for(vecItr=many.begin(),impItr=subImport.begin();
480 vecItr!=many.end();++vecItr,++impItr) {
482 RCP<Tpetra::MultiVector<ST,LO,GO,NT> > destVec = *vecItr;
485 const Tpetra::Map<LO,GO,NT> & globalMap = *(Teuchos::get_extra_data<RCP<Tpetra::Map<LO,GO,NT> > >(destVec,
"globalMap"));
488 GO destLDA = destVec->getStride();
489 GO destSize = destVec->getGlobalLength()*destVec->getNumVectors();
490 std::vector<ST> destArray(destSize);
491 Teuchos::ArrayView<ST> destVals(destArray);
492 destVec->get1dCopy(destVals,destLDA);
493 Tpetra::MultiVector<ST,LO,GO,NT> importVector(rcpFromRef(globalMap),destVals,destLDA,destVec->getNumVectors());
496 importVector.doImport(single,**impItr,Tpetra::INSERT);
498 Tpetra::Import<LO,GO,NT> importer(destVec->getMap(),destVec->getMap());
499 importVector.replaceMap(destVec->getMap());
500 destVec->doImport(importVector,importer,Tpetra::INSERT);