42#ifndef STOKHOS_TPETRA_UTILITIES_UQ_PCE_HPP
43#define STOKHOS_TPETRA_UTILITIES_UQ_PCE_HPP
47#include "Tpetra_Map.hpp"
48#include "Tpetra_MultiVector.hpp"
49#include "Tpetra_CrsGraph.hpp"
50#include "Tpetra_CrsMatrix.hpp"
59 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >
61 const Teuchos::RCP<
const Teuchos::Comm<int> >& comm,
62 const size_t matrix_pce_size) {
64 using Teuchos::arrayView;
66 typedef Kokkos::Compat::KokkosDeviceWrapperNode<Device>
Node;
67 typedef Tpetra::Map<LocalOrdinal,GlobalOrdinal,Node> Map;
68 typedef Tpetra::CrsGraph<LocalOrdinal,GlobalOrdinal,Node> Graph;
72 auto cijk = create_mirror_view(cijk_dev);
73 deep_copy(cijk, cijk_dev);
75 const size_t pce_sz = cijk.dimension();
77 Tpetra::createLocalMapWithNode<LocalOrdinal,GlobalOrdinal,Node>(pce_sz, comm);
79 if (matrix_pce_size == 1) {
80 graph = Tpetra::createCrsGraph(map, 1);
82 for (
size_t i=0; i<pce_sz; ++i) {
84 graph->insertGlobalIndices(row, arrayView(&row, 1));
91 size_t max_num_entry = 0;
92 for (
size_t i=0; i<pce_sz; ++i) {
93 const size_t num_entry = cijk.num_entry(i);
94 max_num_entry = (num_entry > max_num_entry) ? num_entry : max_num_entry;
97 graph = Tpetra::createCrsGraph(map, max_num_entry);
99 for (
size_t i=0; i<pce_sz; ++i) {
101 const size_t num_entry = cijk.num_entry(i);
102 const size_t entry_beg = cijk.entry_begin(i);
103 const size_t entry_end = entry_beg + num_entry;
104 for (
size_t entry = entry_beg; entry < entry_end; ++entry) {
107 graph->insertGlobalIndices(row, arrayView(&
j, 1));
108 graph->insertGlobalIndices(row, arrayView(&k, 1));
112 graph->fillComplete();
123 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >
126 const CijkType& cijk,
127 Teuchos::RCP<
const Tpetra::Map<
LocalOrdinal,
GlobalOrdinal,Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >& flat_domain_map,
128 Teuchos::RCP<
const Tpetra::Map<
LocalOrdinal,
GlobalOrdinal,Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >& flat_range_map,
129 Teuchos::RCP<
const Tpetra::CrsGraph<
LocalOrdinal,
GlobalOrdinal,Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >& cijk_graph,
130 const size_t matrix_pce_size) {
131 using Teuchos::ArrayView;
132 using Teuchos::ArrayRCP;
133 using Teuchos::Array;
137 typedef Kokkos::Compat::KokkosDeviceWrapperNode<Device>
Node;
138 typedef Tpetra::Map<LocalOrdinal,GlobalOrdinal,Node> Map;
139 typedef Tpetra::CrsGraph<LocalOrdinal,GlobalOrdinal,Node> Graph;
144 if (flat_domain_map == Teuchos::null)
145 flat_domain_map =
create_flat_map(*(graph.getDomainMap()), block_size);
148 if (flat_range_map == Teuchos::null)
152 RCP<const Map> flat_col_map =
158 RCP<const Map> flat_row_map;
159 if (graph.getRangeMap() == graph.getRowMap())
160 flat_row_map = flat_range_map;
165 if (cijk_graph == Teuchos::null)
166 cijk_graph = create_cijk_crs_graph<LocalOrdinal,GlobalOrdinal,Device>(
168 flat_domain_map->getComm(),
175 typename Graph::local_inds_host_view_type outer_cols;
176 typename Graph::local_inds_host_view_type inner_cols;
177 size_t max_num_row_entries = graph.getLocalMaxNumRowEntries()*block_size;
178 Array<LocalOrdinal> flat_col_indices;
179 flat_col_indices.reserve(max_num_row_entries);
180 RCP<Graph> flat_graph = rcp(
new Graph(flat_row_map, flat_col_map, max_num_row_entries));
181 const LocalOrdinal num_outer_rows = graph.getLocalNumRows();
182 for (
LocalOrdinal outer_row=0; outer_row < num_outer_rows; outer_row++) {
186 graph.getLocalRowView(outer_row, outer_cols);
190 for (
LocalOrdinal inner_row=0; inner_row < block_size; inner_row++) {
193 const LocalOrdinal flat_row = outer_row*block_size + inner_row;
197 cijk_graph->getLocalRowView(inner_row, inner_cols);
201 flat_col_indices.resize(0);
204 for (
LocalOrdinal outer_entry=0; outer_entry<num_outer_cols;
209 for (
LocalOrdinal inner_entry=0; inner_entry<num_inner_cols;
214 const LocalOrdinal flat_col = outer_col*block_size + inner_col;
215 flat_col_indices.push_back(flat_col);
221 flat_graph->insertLocalIndices(flat_row, flat_col_indices());
226 flat_graph->fillComplete(flat_domain_map, flat_range_map);
235 Teuchos::RCP<
const Tpetra::MultiVector<
typename Storage::value_type,
237 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >
241 Kokkos::Compat::KokkosDeviceWrapperNode<Device> >& vec,
243 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >& flat_map) {
247 typedef typename Storage::value_type BaseScalar;
248 typedef Kokkos::Compat::KokkosDeviceWrapperNode<Device>
Node;
249 typedef Tpetra::MultiVector<BaseScalar,LocalOrdinal,GlobalOrdinal,Node> FlatVector;
250 typedef typename FlatVector::dual_view_type::t_dev flat_view_type;
256 typedef Tpetra::MultiVector<Sacado::UQ::PCE<Storage>,
LocalOrdinal,
GlobalOrdinal, Kokkos::Compat::KokkosDeviceWrapperNode<Device> > mv_type;
257 mv_type& vec_nc =
const_cast<mv_type&
>(vec);
260 flat_view_type flat_vals = vec_nc.getLocalViewDevice(Tpetra::Access::ReadWrite);
263 RCP<FlatVector> flat_vec = rcp(
new FlatVector(flat_map, flat_vals));
272 Teuchos::RCP< Tpetra::MultiVector<
typename Storage::value_type,
274 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >
278 Kokkos::Compat::KokkosDeviceWrapperNode<Device> >& vec,
280 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >& flat_map) {
284 typedef typename Storage::value_type BaseScalar;
285 typedef Kokkos::Compat::KokkosDeviceWrapperNode<Device>
Node;
286 typedef Tpetra::MultiVector<BaseScalar,LocalOrdinal,GlobalOrdinal,Node> FlatVector;
287 typedef typename FlatVector::dual_view_type::t_dev flat_view_type;
290 flat_view_type flat_vals = vec.getLocalViewDevice(Tpetra::Access::ReadWrite);
293 RCP<FlatVector> flat_vec = rcp(
new FlatVector(flat_map, flat_vals));
303 Teuchos::RCP<
const Tpetra::MultiVector<
typename Storage::value_type,
305 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >
309 Kokkos::Compat::KokkosDeviceWrapperNode<Device> >& vec,
311 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >& flat_map) {
312 typedef Kokkos::Compat::KokkosDeviceWrapperNode<Device>
Node;
313 if (flat_map == Teuchos::null) {
318 const Teuchos::RCP< const Tpetra::Map<LocalOrdinal,GlobalOrdinal,Node> > const_flat_map = flat_map;
327 Teuchos::RCP< Tpetra::MultiVector<
typename Storage::value_type,
329 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >
333 Kokkos::Compat::KokkosDeviceWrapperNode<Device> >& vec,
335 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >& flat_map) {
336 typedef Kokkos::Compat::KokkosDeviceWrapperNode<Device>
Node;
337 if (flat_map == Teuchos::null) {
342 const Teuchos::RCP< const Tpetra::Map<LocalOrdinal,GlobalOrdinal,Node> > const_flat_map = flat_map;
350 Teuchos::RCP<
const Tpetra::Vector<
typename Storage::value_type,
352 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >
356 Kokkos::Compat::KokkosDeviceWrapperNode<Device> >& vec_const,
358 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >& flat_map) {
359 typedef Kokkos::Compat::KokkosDeviceWrapperNode<Device>
Node;
361 Teuchos::RCP< Tpetra::MultiVector<typename Storage::value_type,LocalOrdinal,GlobalOrdinal,Node> > flat_mv =
create_flat_vector_view(mv, flat_map);
362 return flat_mv->getVector(0);
370 Teuchos::RCP<
const Tpetra::Vector<
typename Storage::value_type,
372 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >
376 Kokkos::Compat::KokkosDeviceWrapperNode<Device> >& vec,
378 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >& flat_map) {
379 typedef Kokkos::Compat::KokkosDeviceWrapperNode<Device>
Node;
380 if (flat_map == Teuchos::null) {
385 const Teuchos::RCP< const Tpetra::Map<LocalOrdinal,GlobalOrdinal,Node> > const_flat_map = flat_map;
393 Teuchos::RCP< Tpetra::Vector<
typename Storage::value_type,
395 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >
399 Kokkos::Compat::KokkosDeviceWrapperNode<Device> >& vec,
401 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >& flat_map) {
402 typedef Kokkos::Compat::KokkosDeviceWrapperNode<Device>
Node;
404 Teuchos::RCP< Tpetra::MultiVector<typename Storage::value_type,LocalOrdinal,GlobalOrdinal,Node> > flat_mv =
create_flat_vector_view(mv, flat_map);
405 return flat_mv->getVectorNonConst(0);
413 Teuchos::RCP< Tpetra::Vector<
typename Storage::value_type,
415 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >
419 Kokkos::Compat::KokkosDeviceWrapperNode<Device> >& vec,
421 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >& flat_map) {
422 typedef Kokkos::Compat::KokkosDeviceWrapperNode<Device>
Node;
423 if (flat_map == Teuchos::null) {
428 const Teuchos::RCP< const Tpetra::Map<LocalOrdinal,GlobalOrdinal,Node> > const_flat_map = flat_map;
435 typename Device,
typename CijkType>
436 Teuchos::RCP< Tpetra::CrsMatrix<
typename Storage::value_type,
438 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >
442 const Teuchos::RCP<
const Tpetra::CrsGraph<
LocalOrdinal,
GlobalOrdinal,Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >& flat_graph,
443 const Teuchos::RCP<
const Tpetra::CrsGraph<
LocalOrdinal,
GlobalOrdinal,Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >& cijk_graph,
444 const CijkType& cijk_dev) {
445 using Teuchos::ArrayView;
446 using Teuchos::Array;
450 typedef Kokkos::Compat::KokkosDeviceWrapperNode<Device>
Node;
452 typedef typename Storage::value_type BaseScalar;
453 typedef Tpetra::CrsMatrix<Scalar,LocalOrdinal,GlobalOrdinal,Node> Matrix;
454 typedef Tpetra::CrsMatrix<BaseScalar,LocalOrdinal,GlobalOrdinal,Node> FlatMatrix;
458 auto cijk = create_mirror_view(cijk_dev);
459 deep_copy(cijk, cijk_dev);
466 RCP<FlatMatrix> flat_mat = rcp(
new FlatMatrix(flat_graph));
469 typename Matrix::values_host_view_type outer_values;
470 typename Matrix::local_inds_host_view_type outer_cols;
471 typename Matrix::local_inds_host_view_type inner_cols;
472 typename Matrix::local_inds_host_view_type flat_cols;
473 Array<BaseScalar> flat_values;
474 flat_values.reserve(flat_graph->getLocalMaxNumRowEntries());
475 const LocalOrdinal num_outer_rows = mat.getLocalNumRows();
476 for (
LocalOrdinal outer_row=0; outer_row < num_outer_rows; outer_row++) {
479 mat.getLocalRowView(outer_row, outer_cols, outer_values);
483 for (
LocalOrdinal inner_row=0; inner_row < block_size; inner_row++) {
486 const LocalOrdinal flat_row = outer_row*block_size + inner_row;
489 cijk_graph->getLocalRowView(inner_row, inner_cols);
491 ArrayView<const LocalOrdinal> inner_cols_av =
492 Kokkos::Compat::getConstArrayView(inner_cols);
495 const LocalOrdinal num_flat_indices = num_outer_cols*num_inner_cols;
497 flat_values.assign(num_flat_indices, BaseScalar(0));
499 if (matrix_pce_size == 1) {
503 for (
LocalOrdinal outer_entry=0; outer_entry<num_outer_cols;
507 flat_values[outer_entry] = outer_values[outer_entry].coeff(0);
515 const size_t num_entry = cijk.num_entry(inner_row);
516 const size_t entry_beg = cijk.entry_begin(inner_row);
517 const size_t entry_end = entry_beg + num_entry;
518 for (
size_t entry = entry_beg; entry < entry_end; ++entry) {
521 const BaseScalar c = cijk.value(entry);
524 typedef typename ArrayView<const LocalOrdinal>::iterator iterator;
526 std::find(inner_cols_av.begin(), inner_cols_av.end(),
j);
528 std::find(inner_cols_av.begin(), inner_cols_av.end(), k);
529 const LocalOrdinal j_offset = ptr_j - inner_cols_av.begin();
530 const LocalOrdinal k_offset = ptr_k - inner_cols_av.begin();
533 for (
LocalOrdinal outer_entry=0; outer_entry<num_outer_cols;
537 flat_values[outer_entry*num_inner_cols + j_offset] +=
538 c*outer_values[outer_entry].coeff(k);
539 flat_values[outer_entry*num_inner_cols + k_offset] +=
540 c*outer_values[outer_entry].coeff(
j);
549 flat_graph->getLocalRowView(flat_row, flat_cols);
550 flat_mat->replaceLocalValues(flat_row, Kokkos::Compat::getConstArrayView(flat_cols), flat_values());
555 flat_mat->fillComplete(flat_graph->getDomainMap(),
556 flat_graph->getRangeMap());
KokkosClassic::DefaultNode::DefaultNodeType Node
Stokhos::StandardStorage< int, double > Storage
KOKKOS_INLINE_FUNCTION constexpr std::enable_if< is_view_uq_pce< View< T, P... > >::value, unsigned >::type dimension_scalar(const View< T, P... > &view)
Top-level namespace for Stokhos classes and functions.
Teuchos::RCP< Tpetra::CrsGraph< LocalOrdinal, GlobalOrdinal, Kokkos::Compat::KokkosDeviceWrapperNode< Device > > > create_flat_pce_graph(const Tpetra::CrsGraph< LocalOrdinal, GlobalOrdinal, Kokkos::Compat::KokkosDeviceWrapperNode< Device > > &graph, const CijkType &cijk, Teuchos::RCP< const Tpetra::Map< LocalOrdinal, GlobalOrdinal, Kokkos::Compat::KokkosDeviceWrapperNode< Device > > > &flat_domain_map, Teuchos::RCP< const Tpetra::Map< LocalOrdinal, GlobalOrdinal, Kokkos::Compat::KokkosDeviceWrapperNode< Device > > > &flat_range_map, Teuchos::RCP< const Tpetra::CrsGraph< LocalOrdinal, GlobalOrdinal, Kokkos::Compat::KokkosDeviceWrapperNode< Device > > > &cijk_graph, const size_t matrix_pce_size)
Teuchos::RCP< const Tpetra::MultiVector< typename Storage::value_type, LocalOrdinal, GlobalOrdinal, Kokkos::Compat::KokkosDeviceWrapperNode< Device > > > create_flat_vector_view(const Tpetra::MultiVector< Sacado::UQ::PCE< Storage >, LocalOrdinal, GlobalOrdinal, Kokkos::Compat::KokkosDeviceWrapperNode< Device > > &vec, const Teuchos::RCP< const Tpetra::Map< LocalOrdinal, GlobalOrdinal, Kokkos::Compat::KokkosDeviceWrapperNode< Device > > > &flat_map)
Teuchos::RCP< Tpetra::CrsGraph< LocalOrdinal, GlobalOrdinal, Kokkos::Compat::KokkosDeviceWrapperNode< Device > > > create_cijk_crs_graph(const CijkType &cijk_dev, const Teuchos::RCP< const Teuchos::Comm< int > > &comm, const size_t matrix_pce_size)
Teuchos::RCP< Tpetra::CrsMatrix< typename Storage::value_type, LocalOrdinal, GlobalOrdinal, Kokkos::Compat::KokkosDeviceWrapperNode< Device > > > create_flat_matrix(const Tpetra::CrsMatrix< Sacado::UQ::PCE< Storage >, LocalOrdinal, GlobalOrdinal, Kokkos::Compat::KokkosDeviceWrapperNode< Device > > &mat, const Teuchos::RCP< const Tpetra::CrsGraph< LocalOrdinal, GlobalOrdinal, Kokkos::Compat::KokkosDeviceWrapperNode< Device > > > &flat_graph, const Teuchos::RCP< const Tpetra::CrsGraph< LocalOrdinal, GlobalOrdinal, Kokkos::Compat::KokkosDeviceWrapperNode< Device > > > &cijk_graph, const CijkType &cijk_dev)
Teuchos::RCP< Tpetra::Map< LocalOrdinal, GlobalOrdinal, Node > > create_flat_map(const Tpetra::Map< LocalOrdinal, GlobalOrdinal, Node > &map, const LocalOrdinal block_size)