2D Quantum Mechanics: Using Python to Solve Non-Trivial Potentials

Rhett Allain
9 min readNov 18, 2023
Contour plot of a 2D wave function using python.

I’m not really an expert with quantum mechanics calculations — but I need to do some stuff for the class I’m teaching. If I don’t write out my methods, I will forget the next time I teach this course. So, maybe you will find this useful — but future me will certainly like this.

Here’s the situation. Imagine I have a particle on a 2D surface with some type of potential energy function. How can we find the energy eigenvalues and the wave function for something like this? I’m going to show you some python magic — it’s magic because I only partly understand it.

Yes, there are some prerequisites for this method. First, you need to understand how to find analytical solutions for 2D wave functions. We can do this for some situations — like an infinite 2D square well.

You should also be familiar with a numerical method to solve for 1D wave functions. The idea is to break the space into finite elements and then write the Schrödinger equation as a series of linear equations. That can then be represented a system of linear equations represented by a matrix. Here are the details for that.

OK, are you ready? I think it’s important to remember that we are just trying to solve a differential equation. Oh, sure — it’s a differential equation that depends on both time and space. Also, it has complex numbers in there too — but still, it’s just a differential equation with boundary conditions. We can do this.

Wait! One more thing. If I make a mistake or leave something out, you can check out this very awesome paper that describes how to do this stuff.

Wong, Wai Kui. (2021). Solving 2D Time Independent Schrodinger Equation Using Numerical Method. 10.13140/RG.2.2.12835.99360.



Rhett Allain

Physics faculty, science blogger of all things geek. Technical Consultant for CBS MacGyver and MythBusters. WIRED blogger.