Computing the time evolution of wave functions and other complicated physical systems, even with high-end modern hardware, is expensive and time consuming. Neural networks have the ability to act as functional estimators. Finding a way to estimate the Schrödinger equation with neural networks is appealing because it can possibly be done more quickly and efficiently this way.

So what do we practically need to accomplish?

We want to train a neural network to predict wave dynamics.

In order to train a neural network, we need to be able to give it a whole bunch of data. In our case, this data is just the history of the field as the simulation plays out.

So… we must solve the time-dependent Schrödinger equation!

We need to solve for the change in the wave equation some δt later. Computationally, δt must be very small or the calculation will have too much error and become unstable.

This project really isn't concerned with tangible predictions right now, so let's use natural units.

It seems like implementing this should calculate the future state of the field we want.

This project almost exclusively uses PyTorch for compute; it makes matrix math fast and easy anywhere.

Our main simulation class

Starting with the time-dependent Schrödinger equation in natural units:

The formal solution over a time step is:

Using Baker-Campbell-Hausdorff to split the kinetic and potential terms:

This gives us the evolution:

In momentum space, kinetic energy is diagonal:

Transform to position space for potential:

Final momentum space evolution:

This has a lot of advantages compared to the naive method:

But it doesn't handle strong, quickly changing potentials as accurately.

Starting with the time-dependent Schrödinger equation:

The formal solution over timestep δt is:

For small δt, we approximate the exponential using Cayley's form:

Cross multiply to get the evolution equation:

Define time step parameter:

Form matrices:

Solve linear system:

While the direct method is not preferable, it is what I generated the initial dataset with. It can work so long as sufficiently short time steps are used. These simulations each took three hours running at 220 W on an NVIDIA RTX 4090 (0.66 kWh per simulation).

Success! We have data: specifically, 600 frames each from ten 1000×1000 2D simulations. Now we need to make the machines learn.

There are a lot of options when it comes to model architecture, and they will each exhibit unique characteristics.

In a similar way to the methods of simulation, different architectures will be more accurate in certain domains. Some will be more spatially accurate, while others might better capture spectral aspects.

My first approach was to experiment with autoencoders, primarily because they are very simple.

This was not expected to be a very successful approach, as it doesn't inherently do anything special to capture the structure of the system.

It was not successful. The model space consistently collapsed to capture training data, never evolving the system. Maybe this could be blamed on a lack of data.

Fourier neural networks are of particular interest to me because I think they could be incorporated to very effectively capture the spectral parts of a system.

Taking inspiration from the split operator method, I think a model hybrid could be formed which incorporates Fourier networks to handle the momentum terms of the evolution and convolutional networks for the spatial terms.

So far in my experimentation with these architectures, I have not been able to train a numerically stable model.

One thing we know about our desired model is that we want it to evolve locally. The future output of the model should depend only on the immediately adjacent points in the field.

This makes convolutional models particularly interesting.

Convolutional models link locally grouped cells to a single neuron in the next layer. By pooling multiple cells to one in each layer, the dimensionality of the input is reduced.

This makes it a great candidate to convey the local nature of the wave evolution.

This is the first model that made predictions that looked somewhat like time evolution. While the wave packets quickly decohered into noise, the noise still moved as though it had momentum.

The U-Net is a variation of the convolutional autoencoder.

In addition to the standard CAE architecture, the U-Net has:

The hope is that the higher-density, structured skip connections allow the model to capture and predict large temporal differences.

Retrospectively, I think autoencoders are a poor choice for this. The symmetric nature and information bottleneck are not conducive to temporal data, and are better suited for denoising and segmentation.

This has some "success."

But this is still far from an accurate time evolution.

I think the current project still needs a lot of work to be called successful in any way.

I actually started all of this with 2D simulations, which were prohibitively expensive to compute in many cases.

This choice has made it extraordinarily difficult to experiment quickly.

A single epoch of training on ten simulations takes over a day to complete.

In order to better study basic architectural differences, I am going back to 1D simulations for the time being. This is where all the 1D simulations at the beginning came from.

All the original data was generated with the Forward Euler method. It wasn't until I'd shown my professor the original simulations that I was told about Cayley's form.

The results I have seen from Cayley's-form evolution have been very successful, and I would like to generate new datasets based on these.

Currently this project only models single-valued fields. I think an attainable next step would be to simulate and try to model vector fields. If this is done, it is actually more intuitive to go to three rather than two spatial dimensions. Three spatial dimensions would likely be computationally prohibitive with my current methods.

Physics-informed neural networks are an architecture which make calculations on the output based on physical terms like energy conservation, derivatives, symmetries, and so on. These networks allow for better adherence to physical principles, directly enforced through the loss function.

These sorts of networks can actually be trained without any "data."

This project is entirely done in my free time. I am not funded, and all the compute is self-sourced.

If you or anyone you know is interested in facilitating this project on a larger scale, I am always interested in funding, particularly for compute.

At the time of this presentation I was graduating in the spring in Computer Science and Engineering and looking for full-time employment.

Contact from the original presentation: @naeno.com">tsteffen@naeno.com · @uiowa.edu">tstffn@uiowa.edu