47#include "Teko_InterlacedTpetra.hpp"
48#include "Tpetra_Import.hpp"
56namespace TpetraHelpers {
62void 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);
74void 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.getLocalNumElements(),globalMap.getMinGlobalIndex(),
82void 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.getLocalNumElements(),numGlobalVars*sampleMap.getMinGlobalIndex(),vars,comm,subMaps);
100void 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;
165void 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))));
182void 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);
198void 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);
211RCP<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.getLocalNumElements();
244 LO maxNumEntries = A->getGlobalMaxNumRowEntries();
247 auto indices =
typename Tpetra::CrsMatrix<ST,LO,GO,NT>::nonconst_global_inds_host_view_type(Kokkos::ViewAllocateWithoutInitializing(
"rowIndices"),maxNumEntries);
248 auto values =
typename Tpetra::CrsMatrix<ST,LO,GO,NT>::nonconst_values_host_view_type(Kokkos::ViewAllocateWithoutInitializing(
"rowIndices"),maxNumEntries);
251 std::vector<size_t> numEntriesPerRow(numMyRows,0);
253 const size_t invalid = Teuchos::OrdinalTraits<size_t>::invalid();
256 for(LO localRow=0;localRow<numMyRows;localRow++) {
257 size_t numEntries = invalid;
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, indices, 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);
284 numEntriesPerRow[localRow] += numOwnedCols;
287 RCP<Tpetra::CrsMatrix<ST,LO,GO,NT> > mat =
288 rcp(
new Tpetra::CrsMatrix<ST,LO,GO,NT>(rcpFromRef(rowMap), Teuchos::ArrayView<const size_t>(numEntriesPerRow)));
291 std::vector<GO> colIndices(maxNumEntries);
292 std::vector<ST> colValues(maxNumEntries);
296 for(LO localRow=0;localRow<numMyRows;localRow++) {
297 size_t numEntries = invalid;
298 GO globalRow = gRowMap.getGlobalElement(localRow);
299 GO contigRow = rowMap.getGlobalElement(localRow);
301 TEUCHOS_ASSERT(globalRow>=0);
302 TEUCHOS_ASSERT(contigRow>=0);
305 localA.getGlobalRowCopy(globalRow, indices, values, numEntries);
307 for(
size_t localCol=0;localCol<numEntries;localCol++) {
308 GO globalCol = indices(localCol);
311 int block = globalCol / numGlobalVars;
313 bool inFamily =
true;
316 inFamily &= (block*numGlobalVars+colBlockOffset <= globalCol);
317 inFamily &= ((block*numGlobalVars+colBlockOffset+colFamilyCnt) > globalCol);
321 GO familyOffset = globalCol-(block*numGlobalVars+colBlockOffset);
323 colIndices[numOwnedCols] = block*colFamilyCnt + familyOffset;
324 colValues[numOwnedCols] = values(localCol);
331 colIndices.resize(numOwnedCols);
332 colValues.resize(numOwnedCols);
333 mat->insertGlobalValues(contigRow,Teuchos::ArrayView<GO>(colIndices),Teuchos::ArrayView<ST>(colValues));
334 colIndices.resize(maxNumEntries);
335 colValues.resize(maxNumEntries);
339 mat->fillComplete(rcpFromRef(colMap),rcpFromRef(rowMap));
345void 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)
348 int numVarFamily = subMaps.size();
350 TEUCHOS_ASSERT(i>=0 && i<numVarFamily);
351 TEUCHOS_ASSERT(j>=0 && j<numVarFamily);
352 TEUCHOS_ASSERT(mat.isFillComplete());
354 const Tpetra::Map<LO,GO,NT> & gRowMap = *subMaps[i].second;
355 const Tpetra::Map<LO,GO,NT> & rowMap = *Teuchos::get_extra_data<RCP<Tpetra::Map<LO,GO,NT> > >(subMaps[i].second,
"contigMap");
356 const Tpetra::Map<LO,GO,NT> & colMap = *Teuchos::get_extra_data<RCP<Tpetra::Map<LO,GO,NT> > >(subMaps[j].second,
"contigMap");
357 int colFamilyCnt = subMaps[j].first;
361 GO numGlobalVars = 0;
362 GO rowBlockOffset = 0;
363 GO colBlockOffset = 0;
364 for(
int k=0;k<numVarFamily;k++) {
365 numGlobalVars += subMaps[k].first;
368 if(k<i) rowBlockOffset += subMaps[k].first;
369 if(k<j) colBlockOffset += subMaps[k].first;
373 Tpetra::Import<LO,GO,NT>
import(A->getRowMap(),rcpFromRef(gRowMap));
374 Tpetra::CrsMatrix<ST,LO,GO,NT> localA(rcpFromRef(gRowMap),0);
375 localA.doImport(*A,
import,Tpetra::INSERT);
379 mat.setAllToScalar(0.0);
382 LO numMyRows = rowMap.getLocalNumElements();
383 GO maxNumEntries = A->getGlobalMaxNumRowEntries();
386 auto indices =
typename Tpetra::CrsMatrix<ST,LO,GO,NT>::nonconst_global_inds_host_view_type(Kokkos::ViewAllocateWithoutInitializing(
"rowIndices"),maxNumEntries);
387 auto values =
typename Tpetra::CrsMatrix<ST,LO,GO,NT>::nonconst_values_host_view_type(Kokkos::ViewAllocateWithoutInitializing(
"rowIndices"),maxNumEntries);
390 std::vector<GO> colIndices(maxNumEntries);
391 std::vector<ST> colValues(maxNumEntries);
393 const size_t invalid = Teuchos::OrdinalTraits<size_t>::invalid();
397 for(LO localRow=0;localRow<numMyRows;localRow++) {
398 size_t numEntries = invalid;
399 GO globalRow = gRowMap.getGlobalElement(localRow);
400 GO contigRow = rowMap.getGlobalElement(localRow);
402 TEUCHOS_ASSERT(globalRow>=0);
403 TEUCHOS_ASSERT(contigRow>=0);
406 localA.getGlobalRowCopy(globalRow, indices, values, numEntries);
409 for(
size_t localCol=0;localCol<numEntries;localCol++) {
410 GO globalCol = indices(localCol);
413 int block = globalCol / numGlobalVars;
415 bool inFamily =
true;
418 inFamily &= (block*numGlobalVars+colBlockOffset <= globalCol);
419 inFamily &= ((block*numGlobalVars+colBlockOffset+colFamilyCnt) > globalCol);
423 GO familyOffset = globalCol-(block*numGlobalVars+colBlockOffset);
425 colIndices[numOwnedCols] = block*colFamilyCnt + familyOffset;
426 colValues[numOwnedCols] = values(localCol);
433 colIndices.resize(numOwnedCols);
434 colValues.resize(numOwnedCols);
435 mat.sumIntoGlobalValues(contigRow,Teuchos::ArrayView<GO>(colIndices),Teuchos::ArrayView<ST>(colValues));
436 colIndices.resize(maxNumEntries);
437 colValues.resize(maxNumEntries);
439 mat.fillComplete(rcpFromRef(colMap),rcpFromRef(rowMap));
444void many2one(Tpetra::MultiVector<ST,LO,GO,NT> & one,
const std::vector<RCP<
const Tpetra::MultiVector<ST,LO,GO,NT> > > & many,
445 const std::vector<RCP<Tpetra::Export<LO,GO,NT> > > & subExport)
448 std::vector<RCP<const Tpetra::MultiVector<ST,LO,GO,NT> > >::const_iterator vecItr;
449 std::vector<RCP<Tpetra::Export<LO,GO,NT> > >::const_iterator expItr;
452 for(vecItr=many.begin(),expItr=subExport.begin();
453 vecItr!=many.end();++vecItr,++expItr) {
456 RCP<const Tpetra::MultiVector<ST,LO,GO,NT> > srcVec = *vecItr;
459 const Tpetra::Map<LO,GO,NT> & globalMap = *(Teuchos::get_extra_data<RCP<Tpetra::Map<LO,GO,NT> > >(srcVec,
"globalMap"));
462 GO lda = srcVec->getStride();
463 GO srcSize = srcVec->getGlobalLength()*srcVec->getNumVectors();
464 std::vector<ST> srcArray(srcSize);
465 Teuchos::ArrayView<ST> srcVals(srcArray);
466 srcVec->get1dCopy(srcVals,lda);
467 Tpetra::MultiVector<ST,LO,GO,NT> exportVector(rcpFromRef(globalMap),srcVals,lda,srcVec->getNumVectors());
470 one.doExport(exportVector,**expItr,Tpetra::INSERT);
475void one2many(std::vector<RCP<Tpetra::MultiVector<ST,LO,GO,NT> > > & many,
const Tpetra::MultiVector<ST,LO,GO,NT> & single,
476 const std::vector<RCP<Tpetra::Import<LO,GO,NT> > > & subImport)
479 std::vector<RCP<Tpetra::MultiVector<ST,LO,GO,NT> > >::const_iterator vecItr;
480 std::vector<RCP<Tpetra::Import<LO,GO,NT> > >::const_iterator impItr;
483 for(vecItr=many.begin(),impItr=subImport.begin();
484 vecItr!=many.end();++vecItr,++impItr) {
486 RCP<Tpetra::MultiVector<ST,LO,GO,NT> > destVec = *vecItr;
489 const Tpetra::Map<LO,GO,NT> & globalMap = *(Teuchos::get_extra_data<RCP<Tpetra::Map<LO,GO,NT> > >(destVec,
"globalMap"));
492 GO destLDA = destVec->getStride();
493 GO destSize = destVec->getGlobalLength()*destVec->getNumVectors();
494 std::vector<ST> destArray(destSize);
495 Teuchos::ArrayView<ST> destVals(destArray);
496 destVec->get1dCopy(destVals,destLDA);
497 Tpetra::MultiVector<ST,LO,GO,NT> importVector(rcpFromRef(globalMap),destVals,destLDA,destVec->getNumVectors());
500 importVector.doImport(single,**impItr,Tpetra::INSERT);
502 Tpetra::Import<LO,GO,NT> importer(destVec->getMap(),destVec->getMap());
503 importVector.replaceMap(destVec->getMap());
504 destVec->doImport(importVector,importer,Tpetra::INSERT);