print(md.correlation.msd(trajectory[0], trajectory[-1]))

result_msd = md.correlation.shifted_correlation(
    md.correlation.msd, oxygen_water, segments=10, skip=0.1, average=True
)
result_msd_nojump = md.correlation.shifted_correlation(
    md.correlation.msd, oxygen_water.nojump, segments=10, skip=0.1, average=True
)

plt.figure()
plt.plot(result_msd[0], result_msd[1], ".", label="pbc")
plt.plot(result_msd_nojump[0], result_msd_nojump[1], "-", label="nojump")
plt.legend()
plt.xscale("log")
plt.yscale("log")
plt.xlabel(r"$t$ in ps", fontsize=16)
plt.ylabel(r"$\langle r^2\rangle(t)$ in nm$^2$", fontsize=16)
plt.show()
plt.close()