# Load the trajectory
md_trajectory = nlread('alumina_trajectory.hdf5')[-1]

# Calculate the self diffusion
self_diffusion = SelfDiffusion(
    md_trajectory,
    atom_selection=Aluminum
)

# Get the times in ps and the MSD values in Ang**2.
times     = self_diffusion.times()
diffusion = self_diffusion.data()

# Plot the diffusion coefficient calculated with different time limits
model = Plot.PlotModel(ps, cm**2/Second)
model.title().setText('Diffusion coefficient of aluminum')
model.xAxis().setLabel('Time')
model.yAxis().setLabel(r'$D_{self}$')
model.legend().setVisible(False)

# Add diffusion integral
line = Plot.Line(times, diffusion)
line.setColor('blue')

# Add average
average = Plot.Average()
line.addItem(average)
max_time = times[-1]
average.setBounds(max_time*0.1, max_time*0.9)
average.setColor('red')

# Add line to the plot
model.addItem(line)

# Set limits and show the plot
model.setLimits()
Plot.show(model)