KindXiaoming / pykan

Kolmogorov Arnold Networks
MIT License
13.67k stars 1.21k forks source link

Attempting to solve the Time-Independent Schrodinger's Equation in 1D for a particle in a box problem using a KAN model #61

Open bkataru opened 2 months ago

bkataru commented 2 months ago

Hello!

First off, I would like to say, excellent work on KANs and this library. I read your thoughts on the recent explosion in attention that KANs have been receiving and wholeheartedly agree that while KANs may not serve the best purpose in trying to be frontrunners to replace current traditional neural network architectures like MLPs, there is immense potential for their applicability in Physics, especially for symbolic regression tasks and (re)discovering formulae from data.

Recently, me and my advisor had been investigating into deriving novel symbolic series expansions of neutrino oscillation probabilities up to 1st and 2nd orders in oscillation parameters such as the mass-hierarchy parameter and the mixing angles, but we ran into a wall due to the limitations posed by Mathematica and current symbolic CAS systems. Although we do not immediately see how KANs could help us overcome these limitations, we have been considering giving it another shot, this time with the power of KANs by our side.

Coming to my question, I've been playing around with the PDE solving example in the library's documentation and trying to write a similar example for solving the TI Schrodinger's equation in 1D for a classic particle in a box problem and I ran into a roadblock.

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

Generally when we attempt to solve Schrodinger's equation for a particle in a box problem in wave mechanics, we obtain both the energy eigenvalues as well as the solution to the wavefunction in the position basis. If I am to implement a KAN model for this task, would I have to have energy E as a parameter/variable similar to the position x and have the KAN model infer the energy eigenvalue's symbolic relationships by itself, or would I have to specify E by fixing n (the energy level) ahead of time, use the energy eigenvalue equation for this specific system, and utilise that info in the process?

I would really not like to solve for E ahead of time because that just makes it less general and in a setting where one is trying to derive symbolic expressions from data, it is not immediately aware how the energy eigenvalues arise for a given system if there is no further information that can be grokked. Furthermore, if I am to treat energy as a variable, the energy eigenvalues are discrete so I assume this is something I have to encode into the sampling process or would a KAN model be able to identify the discrete nature of this parameter itself if I sample it randomly (which implicitly assumes a continuous nature)?

Is there support for pykan to incorporate and work with parameters that are not variables as of now?

Any help would be greatly appreciated. Cheers!

KindXiaoming commented 2 months ago

Hi, thank you for your kind words! maybe you can treat E as a separate trainable parameter, and use the PINN loss as your loss function. However, you can only get one eigenvalue at one time. But if you try initializing E with different values, you may get multiple solutions. Here is an example for PDE: https://kindxiaoming.github.io/pykan/Examples/Example_6_PDE.html, I guess your only modification is to add E as a trainable parameter (and register it to optimizer).