Don't set positions / velocities to NaNs

Addresses #38
choderalab · May 27, 2017 · e51285f · e51285f
1 parent e20a87b
commit e51285f
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/benchmark/testsystems/bookkeepers.py b/benchmark/testsystems/bookkeepers.py
@@ -248,6 +248,12 @@ def collect_protocol_samples(self, n_protocol_samples, protocol_length, marginal
             else:
                 raise NotImplementedError("`marginal` must be either 'configuration' or 'full'")
 
+            if (np.isnan(x_1).sum() + np.isnan(v_1).sum()) > 0:
+                W_shads_R *= np.nan
+                W_shads_F *= np.nan
+                print("NaNs encountered! Terminating early...")
+                break
+
             W_shads_R[i] = self.accumulate_shadow_work(x_1, v_1, protocol_length)
 
             # if we've encountered any NaNs, terminate early