Add legend function to experiment.
Minor cleanup to plots code.