A python version code basically from @zhaiyusci 's C++ code work.

With wavefunction plot function add by me.

Still ugly because I need learn more about python.