small integrator changes

johnaparker · Apr 21, 2020 · db3dd2e · db3dd2e
1 parent ebba035
commit db3dd2e
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 5 deletions.
diff --git a/stoked/integrators/basic_integrator.py b/stoked/integrators/basic_integrator.py
@@ -7,11 +7,12 @@ def __init__(self):
 
     def bd_step(self):
         F = self.solve_forces()
-        v0 = self.solver.velocity
-        v1 = self.alpha_T*F
+        Fr = self.random_force()
+        F += Fr
 
+        v1 = self.alpha_T*F
         self.solver.velocity = v1
-        self.solver.position += self.dt*v0
+        self.solver.position += self.dt*v1
 
         if self.solver.rotating:
             T = self.solve_torques()
@@ -21,6 +22,9 @@ def bd_step(self):
 
     def ld_step(self):
         F = self.solve_forces()
+        Fr = self.random_force()
+        F += Fr
+
         mass = self.solver.mass
         v0 = self.solver.velocity
         v1 = self.alpha_T*F + (v0 - self.alpha_T*F)*np.exp(-self.dt/mass/self.alpha_T)

diff --git a/stoked/integrators/integrator.py b/stoked/integrators/integrator.py
@@ -27,7 +27,10 @@ def initialize(self, solver):
 
     def solve_forces(self):
         F = self.solver._total_force(self.solver.time, self.solver.position, self.solver.orientation)
-        F += random_force(self.alpha_T, self.solver.temperature, self.dt)
+        return F
+
+    def random_force(self):
+        F = random_force(self.alpha_T, self.solver.temperature, self.dt)
         return F
 
     def solve_torques(self):

diff --git a/stoked/integrators/predictor_corrector_integrator.py b/stoked/integrators/predictor_corrector_integrator.py
@@ -10,6 +10,8 @@ def bd_step(self):
         o0 = np.copy(self.solver.orientation)
 
         F = self.solve_forces()
+        Fr = self.random_force()
+        F += Fr
         v1 = self.alpha_T*F
 
         self.solver.velocity = v1
@@ -30,9 +32,10 @@ def bd_step(self):
         self.pre_step()
 
         F = self.solve_forces()
+        F += Fr
         v2 = self.alpha_T*F
         self.solver.velocity = (v1 + v2)/2
-        self.position = r0 + self.dt*self.solver.velocity
+        self.solver.position = r0 + self.dt*self.solver.velocity
 
         if self.solver.rotating:
             T = self.solve_torques()