improvements to xrb plotting script

[zhichen3]

PR summary

PR motivation

PR checklist

• [ ] test suite needs to be run on this PR
• [ ] this PR will change answers in the test suite to more than roundoff level
• [ ] all newly-added functions have docstrings as per the coding conventions
• [ ] the CHANGES file has been updated, if appropriate
• [ ] if appropriate, this change is described in the docs