Combine force reset and langevin force into single for_each_particle

src/core/forces.cpp

@@ -48,6 +48,7 @@
 #include "short_range_loop.hpp"
 #include "system/System.hpp"
 #include "thermostat.hpp"
+#include "thermostats/langevin_inline.hpp"
 #include "virtual_sites/relative.hpp"
 
 #include <utils/Vector.hpp>
@@ -87,6 +88,47 @@ static ParticleForce external_force(Particle const &p) {
   return f;
 }
 
+/** Combined force initialization and Langevin noise application */
+void init_forces_and_thermostat(const CellStructure &cell_structure,
+                                System::System &system) {
+#ifdef CALIPER
+  CALI_CXX_MARK_FUNCTION;
+#endif
+
+  auto const &propagation = *system.propagation;
+  auto const &thermostat = *system.thermostat;
+  auto const kT = thermostat.kT;
+  auto const time_step = system.get_time_step();
+
+  // Check if Langevin thermostat is active
+  bool const langevin_active =
+      thermostat.langevin &&
+      (propagation.used_propagations &
+       (PropagationMode::TRANS_LANGEVIN | PropagationMode::ROT_LANGEVIN));
+
+  // Single pass over all local particles
+  cell_structure.for_each_local_particle([&](Particle &p) {
+    // Initialize force with external forces (original init_forces logic)
+    p.force_and_torque() = external_force(p);
+
+    // Apply Langevin noise if thermostat is active (original
+    // thermostat_force_init logic)
+    if (langevin_active) {
+      auto const &langevin = *thermostat.langevin;
+      if (propagation.should_propagate_with(p, PropagationMode::TRANS_LANGEVIN))
+        p.force() += friction_thermo_langevin(langevin, p, time_step, kT);
+#ifdef ROTATION
+      if (propagation.should_propagate_with(p, PropagationMode::ROT_LANGEVIN))
+        p.torque() += convert_vector_body_to_space(
+            p, friction_thermo_langevin_rotation(langevin, p, time_step, kT));
+#endif
+    }
+  });
+
+  // Initialize ghost forces (unchanged)
+  init_forces_ghosts(cell_structure);
+}
+
 void init_forces(const CellStructure &cell_structure) {
 #ifdef CALIPER
   CALI_CXX_MARK_FUNCTION;
@@ -150,8 +192,8 @@ void System::System::calculate_forces() {
     npt_inst_pressure->p_vir = Utils::Vector3d{};
   }
 #endif
-  init_forces(*cell_structure);
-  thermostat_force_init();
+  // Use combined function instead of two separate calls
+  init_forces_and_thermostat(*cell_structure, *this);
 
   calc_long_range_forces(particles);
 

src/core/forces.hpp

@@ -34,6 +34,10 @@
 /** Assign external forces/torques to real particles and zero to ghosts. */
 void init_forces(const CellStructure &cell_structure);
 
+/** Combined force initialization and Langevin noise application. */
+void init_forces_and_thermostat(const CellStructure &cell_structure,
+                                class System::System &system);
+
 /** Set forces of all ghosts to zero */
 void init_forces_ghosts(const CellStructure &cell_structure);
 

src/core/integrate.cpp

@@ -317,26 +317,6 @@ static void resort_particles_if_needed(System::System &system) {
   }
 }
 
-void System::System::thermostat_force_init() {
-  auto const &propagation = *this->propagation;
-  if ((not thermostat->langevin) or ((propagation.used_propagations &
-                                      (PropagationMode::TRANS_LANGEVIN |
-                                       PropagationMode::ROT_LANGEVIN)) == 0)) {
-    return;
-  }
-  auto const &langevin = *thermostat->langevin;
-  auto const kT = thermostat->kT;
-  cell_structure->for_each_local_particle([&](Particle &p) {
-    if (propagation.should_propagate_with(p, PropagationMode::TRANS_LANGEVIN))
-      p.force() += friction_thermo_langevin(langevin, p, time_step, kT);
-#ifdef ROTATION
-    if (propagation.should_propagate_with(p, PropagationMode::ROT_LANGEVIN))
-      p.torque() += convert_vector_body_to_space(
-          p, friction_thermo_langevin_rotation(langevin, p, time_step, kT));
-#endif
-  });
-}
-
 /** @brief Calls the hook for propagation kernels before the force calculation
  *  @return whether or not to stop the integration loop early.
  */

src/core/system/System.hpp

@@ -239,8 +239,6 @@ class System : public std::enable_shared_from_this<System> {
   int integrate_with_signal_handler(int n_steps, int reuse_forces,
                                     bool update_accumulators);
 
-  /** @brief Calculate initial particle forces from active thermostats. */
-  void thermostat_force_init();
   /** @brief Calculate particle-lattice interactions. */
   void lb_couple_particles();
 

testsuite/python/caliper.py

@@ -30,14 +30,14 @@
   Initial Force Calculation
     calculate_forces
       copy_particles_to_GPU
-      init_forces
+      init_forces_and_thermost
       calc_long_range_forces
       short_range_loop
       copy_forces_from_GPU
   Integration loop
     calculate_forces
       copy_particles_to_GPU
-      init_forces
+      init_forces_and_thermost
       calc_long_range_forces
       short_range_loop
       copy_forces_from_GPU