From 2ce31aa4e85734d74a7d6c75c87fbabf9aec09c9 Mon Sep 17 00:00:00 2001
From: Sam Reeve <6740307+streeve@users.noreply.github.com>
Date: Fri, 17 Nov 2023 15:01:21 -0500
Subject: [PATCH] Generalize for non-aligned pre-notches

---
 src/CabanaPD_Prenotch.hpp | 60 ++++++++++++++++++++++++++++++++-------
 1 file changed, 50 insertions(+), 10 deletions(-)

diff --git a/src/CabanaPD_Prenotch.hpp b/src/CabanaPD_Prenotch.hpp
index d29aa794..def35bc4 100644
--- a/src/CabanaPD_Prenotch.hpp
+++ b/src/CabanaPD_Prenotch.hpp
@@ -195,22 +195,43 @@ int bondPrenotchIntersection( const Kokkos::Array<double, 3> v1,
     return keep_bond;
 }
 
-template <std::size_t NumNotch>
+template <std::size_t NumNotch, std::size_t NumVector = 1>
 struct Prenotch
 {
     static constexpr std::size_t num_notch = NumNotch;
-    Kokkos::Array<double, 3> _v1;
-    Kokkos::Array<double, 3> _v2;
-    Kokkos::Array<Kokkos::Array<double, 3>, num_notch> _p0_list;
+    static constexpr std::size_t num_vector = NumNotch;
+    Kokkos::Array<Kokkos::Array<double, 3>, num_vector> _v1;
+    Kokkos::Array<Kokkos::Array<double, 3>, num_vector> _v2;
+    Kokkos::Array<Kokkos::Array<double, 3>, num_notch> _p0;
+    bool fixed_orientation;
 
+    // Default constructor
     Prenotch() {}
 
+    // Constructor if all pre-notches are oriented the same way (e.g.
+    // Kalthoff-Winkler).
     Prenotch( Kokkos::Array<double, 3> v1, Kokkos::Array<double, 3> v2,
-              Kokkos::Array<Kokkos::Array<double, 3>, num_notch> p0_list )
+              Kokkos::Array<Kokkos::Array<double, 3>, num_notch> p0 )
+        : _v1( { v1 } )
+        , _v2( { v2 } )
+        , _p0( p0 )
+    {
+        fixed_orientation = true;
+    }
+
+    // Constructor for general case of any orientation for any number of
+    // pre-notches.
+    Prenotch( Kokkos::Array<Kokkos::Array<double, 3>, num_vector> v1,
+              Kokkos::Array<Kokkos::Array<double, 3>, num_vector> v2,
+              Kokkos::Array<Kokkos::Array<double, 3>, num_notch> p0 )
         : _v1( v1 )
         , _v2( v2 )
-        , _p0_list( p0_list )
+        , _p0( p0 )
     {
+        static_assert(
+            num_vector == num_notch,
+            "Number of orientation vectors must match number of pre-notches." );
+        fixed_orientation = false;
     }
 
     template <class ExecSpace, class NeighborView, class Particles,
@@ -221,11 +242,14 @@ struct Prenotch
         auto x = particles.sliceReferencePosition();
         Kokkos::RangePolicy<ExecSpace> policy( 0, particles.n_local );
 
-        auto v1 = _v1;
-        auto v2 = _v2;
-        for ( std::size_t p = 0; p < _p0_list.size(); p++ )
+        for ( std::size_t p = 0; p < _p0.size(); p++ )
         {
-            auto p0 = _p0_list[p];
+            // These will always be different positions.
+            auto p0 = _p0[p];
+            // These may all have the same orientation or all different
+            // orientations.
+            auto v1 = getV1( p );
+            auto v2 = getV2( p );
             auto notch_functor = KOKKOS_LAMBDA( const int i )
             {
                 std::size_t num_neighbors =
@@ -252,6 +276,22 @@ struct Prenotch
             Kokkos::parallel_for( "CabanaPD::Prenotch", policy, notch_functor );
         }
     }
+
+    auto getV1( const int p )
+    {
+        if ( fixed_orientation )
+            return _v1[0];
+        else
+            return _v1[p];
+    }
+
+    auto getV2( const int p )
+    {
+        if ( fixed_orientation )
+            return _v2[0];
+        else
+            return _v2[p];
+    }
 };
 
 } // namespace CabanaPD