Some note on creating a two-dimensional Schrodinger equation solver with neural network.

As a python beginner, an efficient way to learn programming is to solve something that you are interested in. In my case, solving some physics problem. 
Schrodinger equation is a basic equation in quantum physics, which describe the wave properties and the energy eigenstates of the particle. Exact solutions of some simple system in low dimensions are well known. After studying [1], a one dimensional Schrodinger equation solver with the neural network, I would like to do the same thing in two dimensions. I will put some codes on Github recently, so keep following for updating information.

Start from one-dimensional system, the time-independent Schrodinger equation is:

$ -\frac{\hbar^2}{2m}  \frac{d^2}{dx^2} \Psi(x) + V(x) \Psi(x) = E \Psi(x) $

we use gradient descent algorithm in TensorFlow to approximate the ground state solution for training data. Slicing space into discrete points, the neural network here contains two hidden layers and 128 nodes for each layer.

Fig. 1: Structure of neural network.

Mathematically:
$h_{1} = s(W_1 V(x_{i})+b_1)$
$h_{2} = s(W_2 h_1 + b_2) $
$\Psi(x_{i}) = s(W_3 h_2 +b_3)$ 
with $s(k) = \ln(1+ e^k)$
And cost function:
$C(W, b) = \sum\limits_{k = 1}^N \int_{0}^1 (\Psi^{actual}_k(x)-\Psi^{predicted}_k(x))^2 dx$
With $240 \times 40 = 9600$ random generated potential in three different types (step, linear, fourier), we trained the neural network.After 200000 steps, the average training cost is 0.07 and the validation cost is 0.013.  
The result of one of the $9600$ potentials is:

Fig. 2: one-dimensional result.

What about two dimensions? Just look at the time-independent two-dimensional Schrodinger equation first:

$ -\frac{\hbar^2}{2m}  \nabla_{(2)}^2 \Psi(x,y) + V(x,y) \Psi(x,y) = E \Psi(x,y) $

In one-dimensional problem, we sliced 1d space into $128$ discrete points, should we slice 2d space into $128 \times 128 = 16384$ points? 
Yes and no, we need a $128 \times 128$ outputs but only need $128 + 128$ inputs from the fact that the wavefunction is separable:
$\Psi(x, y) = \psi_x(x) \psi_y(y)$
so we can trained $\psi_x(x)$ and $\psi_y(y)$ independently and do a tensorproduct to make it a $128 \times 128$ tensor. One of the potential from $9600 \times 9600$ combinations is 

Fig. 3: potential example.

and corresponding neural network wavefunction and ground state is 

Fig. 4: 2D-neural network wavefunction example.

Fig. 5: ground state solution example.

or in 3d plot:

Fig. 6: 3d plot about the wavefunction. (exciting about learning new function in matplot.)

There are many ways to improve the results, this is just an exercise and note to practice python and neural network method.
Btw, this is the first post on this blog, there will be some other note on a different topic of physics, like the black hole or time machine maybe.
Some other potentials just for fun:

Ref:
[1]one dimensional Schrodinger equation solver with neural network

https://becominghuman.ai/solving-schr%C3%B6dingers-equation-with-deep-learning-f9f6950a7c0e

[2]Python package for calculating wavefunctions for 1D and 2D potentials:

https://github.com/jrjohansson/wavefunction

[3]color bar examples:

https://matplotlib.org/examples/color/colormaps_reference.html

[4]matplot3d tutorial:

https://matplotlib.org/mpl_toolkits/mplot3d/tutorial.html

[5]Deep Learning from scratch by 斎藤康毅 

https://github.com/oreilly-japan/deep-learning-from-scratch