WIP subtract com in diffusion option

cascade/utils.py

@@ -141,7 +141,10 @@ def unwrap_trajectory(traj: list[ase.Atoms]) -> np.ndarray:
     return np.asarray(out)
 
 
-def calculate_sqared_disp(traj: np.ndarray, n_jobs=None, verbose=1) -> np.ndarray:
+def calculate_sqared_disp(traj: np.ndarray, 
+                          n_jobs: int = None, 
+                          verbose: int = 1,
+                          subtract_com_shift: bool = False) -> np.ndarray:
     """Calculate the MSD time correlation function
     
     Calculates the MSD for every length of time accessible in the trajectory, 
@@ -171,6 +174,10 @@ def calc_for_step_size(step_size):
             stop_ix = start_ix + step_size
             # stop - start
             dr = traj[stop_ix, :] - traj[start_ix, :]
+            if subtract_com_shift: 
+                delta_com = traj[stop_ix, :].mean(0) - traj[start_ix, :].mean(0)
+                import pdb; pdb.set_trace()
+                dr -= delta_com
             disp_sq[:, start_ix] = (dr * dr).sum(1)
         # subtract 1 since its zero indexed
         return disp_sq.mean()