Writing scripts in for example Python to produce high-quality figures allows you in a fast and efficient way to produce scientific results that can be included in a report. Furthermore, many operations can easily be automated, avoding thereby tedious repetitions of commands, as well as possible errors. Here we present a simple Python program which solves parts of project 2 for one quantum mechanical particle in a harmonic oscillator potential. The code plots the radial distribution of the three lowest-lying states, in addition to displaying the lowest three eigenvalues. It is easy to modify the trapping potential and run numerical experiments and test different boundary conditions. The plot is obtained using matplotlib, a Python plotting library which produces publication quality figures in a variety of formats and interactive environments across platforms.