pitfall in measuring angle autocorrelation function

angle_wrapping.py

import numpy as np


def wrap(x):
    """from https://github.com/openmm/openmm/issues/2520#issuecomment-575414793"""
    for i in range(len(x)-1):
        if x[i+1]-x[i] > np.pi:
            x[i+1] -= 2*np.pi
        elif x[i+1]-x[i] < -np.pi:
            x[i+1] += 2*np.pi
    minx = x[0]
    maxx = x[0]
    for i in range(len(x)-1):
        if x[i+1] > maxx and x[i+1] > minx+2*np.pi:
            x[i+1] -= 2*np.pi
        elif x[i+1] < minx and x[i+1] < maxx-2*np.pi:
            x[i+1] += 2*np.pi
        minx = min(minx, x[i+1])
        maxx = max(maxx, x[i+1])

        
def random_walk_on_unit_circle(x0, sigma=0.01, n_steps=10000):
    traj = np.zeros((n_steps, 2))
    traj[0] = x0
    for i in range(1, n_steps):
        x = traj[i-1]
        # random perturbation
        dx = np.random.randn(2) * sigma
        
        # project perturbation onto tangent-space of circle at x
        dx_prime = dx - np.dot(x, dx)
    
        # project (x + dx') onto unit circle
        traj[i] = (x + dx_prime) / np.linalg.norm(x + dx_prime)
    return traj


np.random.seed(0)
traj = random_walk_on_unit_circle(x0=np.array([0., 1.]), sigma=0.05, n_steps=10000)

raw_angle = np.arctan2(traj[:,1], traj[:,0])
print("# discontinuous jumps in angle traj:  ", sum(np.abs(np.diff(raw_angle)) > np.pi))

# wrap modifies its argument in-place
wrapped_raw_angle = np.array(raw_angle)
wrap(wrapped_raw_angle)
print("# discontinuous jumps in wrapped traj:", sum(np.abs(np.diff(wrapped_raw_angle)) > np.pi))