NLSE
Non linear Schrödinger equation solver
Install / Use
/learn @Quantum-Optics-LKB/NLSEREADME
NLSE
A package to easily simulate all sorts of non linear Schrödinger equations. It uses a split-step spectral scheme to solve the equations.
Installation
First clone the repository:
git clone https://github.com/Quantum-Optics-LKB/NLSE.git
cd NLSE
Then pip install the package:
pip install .
Basic usage
After installing NLSE, you can simply import one of the solvers and instantiate your problem as follows:
N = 2048 # number of points in solver
n2 = -1.6e-9 # nonlinear index in m^2/W
waist = 2.23e-3 # initial beam waist in m
waist2 = 70e-6 # potential beam waist in m
window = 4*waist # total computational window size in m
power = 1.05 # input optical power in W
Isat = 10e4 # saturation intensity in W/m^2
L = 10e-3 # Length of the medium in m
alpha = 20 # linear losses coefficient in m^-1
backend = "GPU" # whether to run on the GPU or the CPU
simu = NLSE(
alpha, power, window, n2, None, L, NX=N, NY=N, Isat=Isat, backend=backend
)
# Define input field and potential
E_0 = np.exp(-(simu.XX**2 + simu.YY**2) / waist**2)
V = -1e-4 * np.exp(-(simu.XX**2 + simu.YY**2) / waist2**2)
simu.out_field(E_0, L, verbose=True, plot=True, precision="single")
Requirements
Supported platforms
This code has been tested on the three main platforms: Linux, MacOs and Windows. The requirements are in the requirements.txt at the root of the repo.
GPU computing
For optimal speed, this code uses your GPU (graphics card). For this, you need specific libraries. For Nvidia cards, you need a CUDA install. For AMD cards, you need a ROCm install. Of course, you need to update your graphics driver to take full advantage of these. In any case we use CuPy for the Python interface to these libraries.
The cupy dependency is not included in setup.py in order to not break installation on platforms that do not support it !
PyFFTW
If the code does not find Cupy, it will fall back to a CPU based implementation that uses the CPU : PyFFTW. To make the best out of your computer, this library is multithreaded. By default it will use all available threads. If this is not what you want, you can disable this by setting the variable pyfftw.config.NUM_THREADS to a number of your choosing.
On Mac, you first need to install FFTW which can be done by simply using Homebrew brew install fftw. Some users reported this didn't work for them, in this case the next best bet is to build it from source following these instructions: FFTW - Installation and customization.
WARNING : The default flag passed to FFTW for planning is FFTW_PATIENT which means that the first run of the code can take a long time. This information is cached so subsequent runs just have to load the plans, removing this computation time.
PyVkFFT
We found out that PyVkFFT was outperforming CuFFT for our application so the GPU implementation uses this library for optimal performance.
Other than this, the code relies on these libraries :
numba: for best CPU performance on Intel CPU's, withicc_rtpicklenumpyscipymatplotlib
Tests
Tests are included to check functionalities and benchmark performance.
You can run all tests by using pytest at the root of the repo.
It will test both CPU and GPU backends (if available).
This can take some time !
The benchmarks can be run using examples/benchmarks.py and compare a "naive" numpy implementation of the main solver loop to our solver.
We also compare for the example of the vortex precession presented in FourierGPE.jl to our solver.
On a Nvidia RTX4090 GPU and Ryzen 7950X CPU, we test our solver to the following results:
How does it work ?
Physical situation
The code offers to solve a typical non linear Schrödinger / Gross-Pitaevskii equation of the type : $$i\partial_{t}\psi = -\frac{1}{2}\nabla^2\psi+V\psi+g|\psi|^2\psi$$
In this particular instance, we solve in the formalism of the propagation of a pulse of light in a non linear medium. Within the paraxial approximation, the propagation equation for the field $E$ in V/m solved is:
$$ i\partial_{z}E = -\frac{1}{2k_0}\nabla_{\perp}^2 E + \frac{D_0}{2}\partial^2_t E -\frac{k_0}{2}\delta n(r) E - n_2 \frac{k_0}{2n}c\epsilon_0|E|^2E $$
Here, the constants are defined as followed :
- $k_0$ : is the electric field wavenumber in $m^{-1}$
- $D_0$ : is the group velocity dispersion (GVD) in $s^2/m$
- $\delta n(\mathbf{r})$ : the "potential" i.e a local change in linear index of refraction. Dimensionless.
- $n_2$ : the non linear index of refraction in $m^2/W$.
- $n$ is the linear index of refraction. In our case 1.
- $c,\epsilon_0$ : the speed of light and electric permittivity of vacuum.
In all generality, the interaction term can be non-local i.e $n_2=n_2(\mathbf{r})$. This means usually that the response will be described as a convolution by some non-local kernel:
$$ n_2(\mathbf{r})|E|^2(\mathbf{r})=n_2\int_{\mathbb{R}^2}\mathrm{d}\mathbf{r}' K(\mathbf{r}-\mathbf{r}')|E|^2(\mathbf{r}'), $$
where $K(\mathbf{r})$ is the non-local kernel, typically the Green function of some diffusion equation.
Please note that all of the code works with the "God given" units i.e SI units !
The NLSE class
The NLSE class aims at providing a minimal yet functional toolbox to solve non-linear Schrödinger type equation in optics / atomic physics settings such as the propagation of light in a Kerr medium or solving the Gross Pitaevskii equation for the evolution of cold gases.
The propagation equation is:
$$ i\partial_{z}E = -\frac{1}{2k_0}\nabla_{\perp}^2 E + -\frac{k_0}{2}\delta n(r) E - n_2 \frac{k_0}{2n}c\epsilon_0|E|^2E $$
Initialization
The physical parameters listed above are defined at the instantiation of the NLSE class (__init__ function).
The backend (GPU or CPU) is tested when the library is imported, but you can then dynamically switch it when instantiating a NLSE class by setting the self.backend attribute to "GPU" or "CPU".
Broadcasting
Since numpy / cupy allow for natural broadcasting of arrays of compatible size, one can leverage this in order to run parallel realizations. For instance, if we wish to propagate various initial state with the same physical parameters,
we simply have to initialize a tensor of fields of dimensions (N_real, Ny, Nx) where N_real is the number of initial states we wish to propagate.
Similarly, one can broadcast over the physical parameters by setting some parameters to be tensors as well. If we wish for instance to study the effect of the variation of $n_2$, one can set the n2 attribute to be a (N_real, 1, 1) tensor.
The field should then be initialized to a (N_real, Ny, Nx) tensor of identical fields and each slice over the first dimension will represent the same field propagated with different parameters.
Finally, one can combine broadcasting over several parameters at the same time: if we wish to do a grid search over $n_2$ and $\alpha$, one can instantiate n2 to be a (N_n2, 1, 1, 1) array, alpha to be a (1, N_alpha, 1, 1) and the field
a (N_n2, N_alpha, Ny, Nx) array.
The take-home message is that the array shape should be compliant with numpy broadcasting rules.
WARNING : For now it is only available on the GPU backend !
Numerical precision
In order to reach the best performance, the numerical precision is hardcoded as a constant variable at the top of nlse.py.
When importing the solvers like NLSE, the data types of the input arrays must match the data type given as input of out_field.
Callbacks
The out_field functions support callbacks with the following signature callback(self, A, z, i) where self is the class instance, A is the field, z is the current position and i the main loop index.
For example if you want to print the number of steps every 100 steps, this is the callback you could write :
def callback(nlse, A, z, i):
n = int(z/nlse.delta_z)
if n % 100 == 0:
print(n)
Notice that since the class instance is passed to the callback, you have access to all of the classes attributes.
Be mindful however that since the callback is running in the main solver loop, this function should not be called too often in order to not slow down the execution too much.
You can find several generic callbacks in the callbacks sublibrary.
Propagation
The out_field method is the main function of the code that propagates the field for an arbitrary distance from an initial state E_in from z=0 (assumed to be the begining of the non linear medium) up to a specified distance z. This function simply works by iterating the spectral solver scheme i.e :
- (If precision is
'double'apply real space terms) - Fourier transforming the field
- Applying the laplacian op
