get view of particle vector/matrix

ECP-copa · Oct 23, 2024 · 2d68f11 · 2d68f11
1 parent 86b41a2
commit 2d68f11
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 0 deletions.
diff --git a/core/src/Cabana_Fields.hpp b/core/src/Cabana_Fields.hpp
@@ -16,6 +16,8 @@
 #ifndef CABANA_FIELDS_HPP
 #define CABANA_FIELDS_HPP
 
+#include <Cabana_BatchedLinearAlgebra.hpp>
+
 namespace Cabana
 {
 namespace Field
@@ -47,6 +49,8 @@ struct Scalar
     static constexpr int size = 1;
     //! Scalar type.
     using data_type = value_type;
+    //! Linear algebra type.
+    using linear_algebra_type = value_type;
 };
 
 //! Vector (1D) particle field type.
@@ -63,6 +67,8 @@ struct Vector
     static constexpr int dim0 = D;
     //! Scalar type.
     using data_type = value_type[D];
+    //! Linear algebra type.
+    using linear_algebra_type = Cabana::LinearAlgebra::VectorView<T, D>;
 };
 
 //! Matrix (2D) particle field type.
@@ -81,6 +87,8 @@ struct Matrix
     static constexpr int dim1 = D1;
     //! Scalar type.
     using data_type = value_type[D0][D1];
+    //! Linear algebra type.
+    using linear_algebra_type = LinearAlgebra::MatrixView<T, D0, D1>;
 };
 
 //---------------------------------------------------------------------------//

diff --git a/core/src/Cabana_ParticleList.hpp b/core/src/Cabana_ParticleList.hpp
@@ -172,6 +172,59 @@ get( ParticleView<VectorLength, FieldTags...>& particle, FieldTag,
         particle.soa(), particle.vectorIndex(), indices... );
 }
 
+//---------------------------------------------------------------------------//
+// Get a view of a particle member as a vector. (Works for both Particle
+// and ParticleView)
+template <class ParticleType, class FieldTag>
+KOKKOS_FORCEINLINE_FUNCTION typename std::enable_if<
+    LinearAlgebra::is_vector<typename FieldTag::linear_algebra_type>::value,
+    typename FieldTag::linear_algebra_type>::type
+get( ParticleType& particle, FieldTag tag )
+{
+    return typename FieldTag::linear_algebra_type(
+        &( Cabana::get( particle, tag, 0 ) ), ParticleType::vector_length );
+}
+
+template <class ParticleType, class FieldTag>
+KOKKOS_FORCEINLINE_FUNCTION typename std::enable_if<
+    LinearAlgebra::is_vector<typename FieldTag::linear_algebra_type>::value,
+    const typename FieldTag::linear_algebra_type>::type
+get( const ParticleType& particle, FieldTag tag )
+{
+    return typename FieldTag::linear_algebra_type(
+        const_cast<typename FieldTag::value_type*>(
+            &( Cabana::get( particle, tag, 0 ) ) ),
+        ParticleType::vector_length );
+}
+
+//---------------------------------------------------------------------------//
+// Get a view of a particle member as a matrix. (Works for both Particle
+// and ParticleView)
+template <class ParticleType, class FieldTag>
+KOKKOS_FORCEINLINE_FUNCTION typename std::enable_if<
+    LinearAlgebra::is_matrix<typename FieldTag::linear_algebra_type>::value,
+    typename FieldTag::linear_algebra_type>::type
+get( ParticleType& particle, FieldTag tag )
+{
+    return typename FieldTag::linear_algebra_type(
+        &( Cabana::get( particle, tag, 0, 0 ) ),
+        ParticleType::vector_length * FieldTag::dim1,
+        ParticleType::vector_length );
+}
+
+template <class ParticleType, class FieldTag>
+KOKKOS_FORCEINLINE_FUNCTION typename std::enable_if<
+    LinearAlgebra::is_matrix<typename FieldTag::linear_algebra_type>::value,
+    const typename FieldTag::linear_algebra_type>::type
+get( const ParticleType& particle, FieldTag tag )
+{
+    return typename FieldTag::linear_algebra_type(
+        const_cast<typename FieldTag::value_type*>(
+            &( Cabana::get( particle, tag, 0, 0 ) ) ),
+        ParticleType::vector_length * FieldTag::dim1,
+        ParticleType::vector_length );
+}
+
 //---------------------------------------------------------------------------//
 //! List of particle fields stored in AoSoA.
 template <class MemorySpace, int VectorLength, class... FieldTags>