import pylab

# Read in the existing trajectory to ge the data
trajectory_file = "Production_Trajectory.hdf5"
trajectory = nlread(trajectory_file, MDTrajectory)[-1]
time = trajectory.times()
pressure = trajectory.pressures()

# Set the average and target pressures
p_tot_avg = numpy.cumsum(pressure.inUnitsOf(GPa)) * GPa / (numpy.arange(pressure.shape[0])+1)
p_target = numpy.ones(pressure.shape[0]) * Pa

# Plot the resulting pressure and average pressure
pylab.figure()
pylab.plot(time.inUnitsOf(picosecond), pressure.inUnitsOf(GPa), label='Pressure (GPa)', linewidth=0.5)
pylab.plot(time.inUnitsOf(picosecond), p_tot_avg.inUnitsOf(GPa), label='P Average', linewidth=2.0)
pylab.plot(time.inUnitsOf(picosecond), p_target.inUnitsOf(GPa), label='P Target', linewidth=2.0)
pylab.xlabel('Time (ps)')
pylab.ylabel('Pressure (GPa)')
pylab.legend()
pylab.show()