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);
242 RCP<Tpetra::CrsMatrix<ST,LO,GO,NT> > mat = Tpetra::createCrsMatrix<ST,LO,GO,NT>(rcpFromRef(rowMap),0);
245 LO numMyRows = rowMap.getNodeNumElements();
246 LO maxNumEntries = A->getGlobalMaxNumRowEntries();
249 std::vector<GO> indices(maxNumEntries);
250 std::vector<ST> values(maxNumEntries);
253 std::vector<GO> colIndices(maxNumEntries);
254 std::vector<ST> colValues(maxNumEntries);
258 for(LO localRow=0;localRow<numMyRows;localRow++) {
259 size_t numEntries = -1;
260 GO globalRow = gRowMap.getGlobalElement(localRow);
261 GO contigRow = rowMap.getGlobalElement(localRow);
263 TEUCHOS_ASSERT(globalRow>=0);
264 TEUCHOS_ASSERT(contigRow>=0);
267 localA.getGlobalRowCopy(globalRow, Teuchos::ArrayView<GO>(indices), Teuchos::ArrayView<ST>(values), numEntries);
269 for(
size_t localCol=0;localCol<numEntries;localCol++) {
270 GO globalCol = indices[localCol];
273 int block = globalCol / numGlobalVars;
275 bool inFamily =
true;
278 inFamily &= (block*numGlobalVars+colBlockOffset <= globalCol);
279 inFamily &= ((block*numGlobalVars+colBlockOffset+colFamilyCnt) > globalCol);
283 GO familyOffset = globalCol-(block*numGlobalVars+colBlockOffset);
285 colIndices[numOwnedCols] = block*colFamilyCnt + familyOffset;
286 colValues[numOwnedCols] = values[localCol];
293 colIndices.resize(numOwnedCols);
294 colValues.resize(numOwnedCols);
295 mat->insertGlobalValues(contigRow,Teuchos::ArrayView<GO>(colIndices),Teuchos::ArrayView<ST>(colValues));
296 colIndices.resize(maxNumEntries);
297 colValues.resize(maxNumEntries);
301 mat->fillComplete(rcpFromRef(colMap),rcpFromRef(rowMap));
307 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)
310 int numVarFamily = subMaps.size();
312 TEUCHOS_ASSERT(i>=0 && i<numVarFamily);
313 TEUCHOS_ASSERT(j>=0 && j<numVarFamily);
314 TEUCHOS_ASSERT(mat.isFillComplete());
316 const Tpetra::Map<LO,GO,NT> & gRowMap = *subMaps[i].second;
317 const Tpetra::Map<LO,GO,NT> & rowMap = *Teuchos::get_extra_data<RCP<Tpetra::Map<LO,GO,NT> > >(subMaps[i].second,
"contigMap");
318 const Tpetra::Map<LO,GO,NT> & colMap = *Teuchos::get_extra_data<RCP<Tpetra::Map<LO,GO,NT> > >(subMaps[j].second,
"contigMap");
319 int colFamilyCnt = subMaps[j].first;
323 GO numGlobalVars = 0;
324 GO rowBlockOffset = 0;
325 GO colBlockOffset = 0;
326 for(
int k=0;k<numVarFamily;k++) {
327 numGlobalVars += subMaps[k].first;
330 if(k<i) rowBlockOffset += subMaps[k].first;
331 if(k<j) colBlockOffset += subMaps[k].first;
335 Tpetra::Import<LO,GO,NT>
import(A->getRowMap(),rcpFromRef(gRowMap));
336 Tpetra::CrsMatrix<ST,LO,GO,NT> localA(rcpFromRef(gRowMap),0);
337 localA.doImport(*A,
import,Tpetra::INSERT);
341 mat.setAllToScalar(0.0);
344 LO numMyRows = rowMap.getNodeNumElements();
345 GO maxNumEntries = A->getGlobalMaxNumRowEntries();
348 std::vector<GO> indices(maxNumEntries);
349 std::vector<ST> values(maxNumEntries);
352 std::vector<GO> colIndices(maxNumEntries);
353 std::vector<ST> colValues(maxNumEntries);
357 for(LO localRow=0;localRow<numMyRows;localRow++) {
358 size_t numEntries = -1;
359 GO globalRow = gRowMap.getGlobalElement(localRow);
360 GO contigRow = rowMap.getGlobalElement(localRow);
362 TEUCHOS_ASSERT(globalRow>=0);
363 TEUCHOS_ASSERT(contigRow>=0);
366 localA.getGlobalRowCopy(globalRow, Teuchos::ArrayView<GO>(indices), Teuchos::ArrayView<ST>(values), numEntries);
369 for(
size_t localCol=0;localCol<numEntries;localCol++) {
370 GO globalCol = indices[localCol];
373 int block = globalCol / numGlobalVars;
375 bool inFamily =
true;
378 inFamily &= (block*numGlobalVars+colBlockOffset <= globalCol);
379 inFamily &= ((block*numGlobalVars+colBlockOffset+colFamilyCnt) > globalCol);
383 GO familyOffset = globalCol-(block*numGlobalVars+colBlockOffset);
385 colIndices[numOwnedCols] = block*colFamilyCnt + familyOffset;
386 colValues[numOwnedCols] = values[localCol];
393 colIndices.resize(numOwnedCols);
394 colValues.resize(numOwnedCols);
395 mat.sumIntoGlobalValues(contigRow,Teuchos::ArrayView<GO>(colIndices),Teuchos::ArrayView<ST>(colValues));
396 colIndices.resize(maxNumEntries);
397 colValues.resize(maxNumEntries);
399 mat.fillComplete(rcpFromRef(colMap),rcpFromRef(rowMap));
404 void many2one(Tpetra::MultiVector<ST,LO,GO,NT> & one,
const std::vector<RCP<
const Tpetra::MultiVector<ST,LO,GO,NT> > > & many,
405 const std::vector<RCP<Tpetra::Export<LO,GO,NT> > > & subExport)
408 std::vector<RCP<const Tpetra::MultiVector<ST,LO,GO,NT> > >::const_iterator vecItr;
409 std::vector<RCP<Tpetra::Export<LO,GO,NT> > >::const_iterator expItr;
412 for(vecItr=many.begin(),expItr=subExport.begin();
413 vecItr!=many.end();++vecItr,++expItr) {
416 RCP<const Tpetra::MultiVector<ST,LO,GO,NT> > srcVec = *vecItr;
419 const Tpetra::Map<LO,GO,NT> & globalMap = *(Teuchos::get_extra_data<RCP<Tpetra::Map<LO,GO,NT> > >(srcVec,
"globalMap"));
422 GO lda = srcVec->getStride();
423 GO srcSize = srcVec->getGlobalLength()*srcVec->getNumVectors();
424 std::vector<ST> srcArray(srcSize);
425 Teuchos::ArrayView<ST> srcVals(srcArray);
426 srcVec->get1dCopy(srcVals,lda);
427 Tpetra::MultiVector<ST,LO,GO,NT> exportVector(rcpFromRef(globalMap),srcVals,lda,srcVec->getNumVectors());
430 one.doExport(exportVector,**expItr,Tpetra::INSERT);
435 void one2many(std::vector<RCP<Tpetra::MultiVector<ST,LO,GO,NT> > > & many,
const Tpetra::MultiVector<ST,LO,GO,NT> & single,
436 const std::vector<RCP<Tpetra::Import<LO,GO,NT> > > & subImport)
439 std::vector<RCP<Tpetra::MultiVector<ST,LO,GO,NT> > >::const_iterator vecItr;
440 std::vector<RCP<Tpetra::Import<LO,GO,NT> > >::const_iterator impItr;
443 for(vecItr=many.begin(),impItr=subImport.begin();
444 vecItr!=many.end();++vecItr,++impItr) {
446 RCP<Tpetra::MultiVector<ST,LO,GO,NT> > destVec = *vecItr;
449 const Tpetra::Map<LO,GO,NT> & globalMap = *(Teuchos::get_extra_data<RCP<Tpetra::Map<LO,GO,NT> > >(destVec,
"globalMap"));
452 GO destLDA = destVec->getStride();
453 GO destSize = destVec->getGlobalLength()*destVec->getNumVectors();
454 std::vector<ST> destArray(destSize);
455 Teuchos::ArrayView<ST> destVals(destArray);
456 destVec->get1dCopy(destVals,destLDA);
457 Tpetra::MultiVector<ST,LO,GO,NT> importVector(rcpFromRef(globalMap),destVals,destLDA,destVec->getNumVectors());
460 importVector.doImport(single,**impItr,Tpetra::INSERT);
462 Tpetra::Import<LO,GO,NT> importer(destVec->getMap(),destVec->getMap());
463 importVector.replaceMap(destVec->getMap());
464 destVec->doImport(importVector,importer,Tpetra::INSERT);