# -*- coding: utf-8 -*-
# Copyright 2012-2013 Pierre de Buyl
# Copyright 2013 Felix Hoëfling
#
# This file is part of pyh5md
#
# pyh5md is free software and is licensed under the modified BSD license (see
# LICENSE file).

import numpy as np
import matplotlib.pyplot as plt
import pyh5md

# Open a H5MD file
f = pyh5md.H5MD_File('walk_1d.h5', 'r')
f.check()

# Open a trajectory group
part = f.particles_group('particles')

# Open trajectory position data element in the trajectory group
part_pos = part.trajectory('position')

# Get data and time
r = part_pos.value
r_time = part_pos.time

# Compute the time-averaged mean-square displacement,
# drop large correlation times due to insufficient statistics
T = r.shape[0]
msd = np.empty((T/4, r.shape[1]))
time = r_time[:T/4]
for n in range(T/4):
    # the sum over "axis=2" is over the spatial components of the positions
    msd[n] = np.mean(np.sum(pow(r[n:] - r[:T-n], 2), axis=2), axis=0)

# Compute mean and standard error of mean (particle- and component-wise)
msd_mean = msd.mean(axis=1)
msd_err = msd.std(axis=1) / np.sqrt(msd.shape[1] - 1)

# Display the MSD and its standard error
plt.plot(time, msd_mean, 'k-', label=r'$\langle [{\bf r}(t)-{\bf r}(0)]^2\rangle$')
plt.plot(time, msd_mean+msd_err, 'k:', label=r'$\langle [{\bf r}(t)-{\bf r}(0)]^2\rangle \pm \sigma$')
plt.plot(time, msd_mean-msd_err, 'k:')
# display reference line for long-time diffusion with D = <a^2> / (2 d <\tau>),
# here: <a^2> = 1, <\tau> = 0.1, and d=1
plt.plot(time, 2 * (.5 * 1 / 0.1 / 1) * time, 'k--', label=r'$2 D t$')
plt.xlabel(r'$t$')
plt.ylabel(r'$MSD(t)$')
plt.xscale('log')
plt.yscale('log')
plt.legend(loc='upper left')

# Create a new figure
plt.figure()

# Obtain and plot the center_of_mass observable
obs_com = f.observable('center_of_mass')
plt.plot(obs_com.time, obs_com.value, 'k-')
plt.xlabel(r'$t$')
plt.ylabel(r'center of mass')

# Close the file
f.close()

plt.show()
