allows particle systems with 7 variables instead of 6

cpp/particles/particles_inner_computer_2nd_order.hpp

@@ -29,6 +29,16 @@
 #include <cstring>
 #include <cassert>
 
+/** \brief Computes action of Stokes drag on particles.
+ *
+ * Either 6 state variables (positions + momenta) or 7 state variables (positions + momenta + value of drag coefficient).
+ * The case of particles with individual values of the drag coefficient is relevant for the case of plastic collisions,
+ * when we do not have a priori knowledge of what the mass/size/drag coefficient of the particle is.
+ *
+ * `compute_interaction_with_extra` is currently only defined for 6 or 7 state variables.
+ *
+ */
+
 template <class real_number, class partsize_t>
 class particles_inner_computer_2nd_order_Stokes{
     double drag_coefficient;
@@ -51,9 +61,10 @@ public:
             real_number particle_state[],
             real_number particle_rhs[],
             const real_number sampled_velocity[]) const{
-        assert(size_particle_state == 6);
-        assert(size_particle_rhs == 6);
         assert(size_particle_rhs_extra == 3);
+        assert(size_particle_state == size_particle_rhs);
+        assert((size_particle_state == 6) || (size_particle_state == 7));
+        if (size_particle_state == 6){
             #pragma omp parallel for
             for(partsize_t idx_part = 0 ; idx_part < number_of_particles ; ++idx_part){
                 particle_rhs[idx_part*size_particle_rhs + IDXC_X] = particle_state[idx_part*size_particle_state + 3 + IDXC_X];
@@ -64,6 +75,19 @@ public:
                 particle_rhs[idx_part*size_particle_rhs + 3 + IDXC_Z] = - this->drag_coefficient * (particle_state[idx_part*size_particle_state + 3 + IDXC_Z] - sampled_velocity[idx_part*size_particle_rhs_extra + IDXC_Z]);
             }
         }
+        else if (size_particle_state == 7){
+            // particle stores its own drag coefficient in the 7th state variable
+            #pragma omp parallel for
+            for(partsize_t idx_part = 0 ; idx_part < number_of_particles ; ++idx_part){
+                particle_rhs[idx_part*size_particle_rhs + IDXC_X] = particle_state[idx_part*size_particle_state + 3 + IDXC_X];
+                particle_rhs[idx_part*size_particle_rhs + IDXC_Y] = particle_state[idx_part*size_particle_state + 3 + IDXC_Y];
+                particle_rhs[idx_part*size_particle_rhs + IDXC_Z] = particle_state[idx_part*size_particle_state + 3 + IDXC_Z];
+                particle_rhs[idx_part*size_particle_rhs + 3 + IDXC_X] = - particle_state[idx_part*size_particle_state + 6] * (particle_state[idx_part*size_particle_state + 3 + IDXC_X] - sampled_velocity[idx_part*size_particle_rhs_extra + IDXC_X]);
+                particle_rhs[idx_part*size_particle_rhs + 3 + IDXC_Y] = - particle_state[idx_part*size_particle_state + 6] * (particle_state[idx_part*size_particle_state + 3 + IDXC_Y] - sampled_velocity[idx_part*size_particle_rhs_extra + IDXC_Y]);
+                particle_rhs[idx_part*size_particle_rhs + 3 + IDXC_Z] = - particle_state[idx_part*size_particle_state + 6] * (particle_state[idx_part*size_particle_state + 3 + IDXC_Z] - sampled_velocity[idx_part*size_particle_rhs_extra + IDXC_Z]);
+            }
+        }
+    }
 
     constexpr static bool isEnable() {
         return true;

cpp/particles/particles_system.hpp

@@ -314,8 +314,8 @@ public:
     }
 
     void complete2ndOrderLoop(const real_number dt) final {
-        assert(size_particle_positions == 6);
-        assert(size_particle_rhs == 6);
+        assert((size_particle_positions == 6) || (size_particle_positions == 7));
+        assert(size_particle_rhs == size_particle_positions);
         std::unique_ptr<real_number[]> sampled_velocity(new real_number[getLocalNbParticles()*3]());
         std::fill_n(sampled_velocity.get(), 3*getLocalNbParticles(), 0);
         sample_compute_field(default_field, sampled_velocity.get());