Fourier-Transformer for solving Fokker-Planck PDE

Transformer for solving the problem in Stochastic Local Volatility

• Basic Component
  • All the files can be found here
  • The problem is denfined in Stochastic Local Volatility, with SDE of 2 dimensions in space and 1 dimension in time.
• Stages of the lab
  • Experiment_0, sampling of the FokkerPlanck PDE, with the help of Quantlib.
  • Experiment_1, FNO in solving the FokkerPlanck PDE.
  • Experiment_2, FNO Transformer in solving the FokkerPlanck PDE.
• Basic Introduction
  • There are also some RL based methods to solve it
    • In solving HJB PDE, Actor-Critic Method for High Dimensional Static Hamilton-Jacobi-Bellman Partial Differential Equations based on Neural Networks
    • In solving Elliptic PDE, Neural Q-learning for solving elliptic PDEs

Week4

2022.May.09, 11:00AM~12:15Noon, with Dr. Shuhao Cao, Ziheng Wang, and Deqing Jiang at Mathematical Institute.

• The interface bwtween Quantlib and Transform
  • Sampling point format
  • Pre-process data format on our own
• The result of the PDE is in a stochastic way, which is totally different from the
• 
  In order to get v, the p is what we want to solve, the transition density. It can be discribed by The Fokker-Planck Equation, which is 2 dimensions in space and 1 dimension in time.
  
• #TODO Problems still remains
  • How to define the input data format
  • How to define the final result of our model
  • Is there any advantage compared with the tradition method Monte Carlo?
• Reference paper
  • Fourier Neural Operator Fourier Neural Operator for Parametric Partial Differential Equations
  • Galerkin Transformer Choose a Transformer Fourier or Galerkin
  • Fourier Transformer Adaptive Fourier Neural Operators: Efficient Token Mixers for Transformers
  • DeepXDE A deep learning library for solving differential equations Documentation

Week3

• Problems remian in this week
  • Only a set of PDEs are suitable for Fourier Transform.
  • Why sample is 256*256 and downsample is 64*64, what has downsampling done here?
  • How MCMC using FNO？
  • To solve different PDEs, do we need to change the code in the FNO?
    • Or the info has been included in the dataset by sampling.
  • We need to solve a two dimensional difussion equation.

    • Fourier-Transformer/fourier_neural_operator/fourier_3d.py

      Lines 1 to 7 in 6fd306e

      	"""
      	@author: Zongyi Li
      	This file is the Fourier Neural Operator for 3D problem
      	such as the Navier-Stokes equation discussed in Section 5.3
      	in the [paper](https://arxiv.org/pdf/2010.08895.pdf),
      	which takes the 2D spatial + 1D temporal equation directly as a 3D problem
      	"""

    • In the class SpectralConv3d(nn.Module):
      

      Fourier-Transformer/fourier_neural_operator/fourier_3d.py

      Lines 61 to 76 in 90686ea

      	#Compute Fourier coeffcients up to factor of e^(- something constant)
      	x_ft = torch.fft.rfftn(x, dim=[-3,-2,-1])
      	# Multiply relevant Fourier modes
      	out_ft = torch.zeros(batchsize, self.out_channels, x.size(-3), x.size(-2), x.size(-1)//2 + 1, dtype=torch.cfloat, device=x.device)
      	out_ft[:, :, :self.modes1, :self.modes2, :self.modes3] = \
      	self.compl_mul3d(x_ft[:, :, :self.modes1, :self.modes2, :self.modes3], self.weights1)
      	out_ft[:, :, -self.modes1:, :self.modes2, :self.modes3] = \
      	self.compl_mul3d(x_ft[:, :, -self.modes1:, :self.modes2, :self.modes3], self.weights2)
      	out_ft[:, :, :self.modes1, -self.modes2:, :self.modes3] = \
      	self.compl_mul3d(x_ft[:, :, :self.modes1, -self.modes2:, :self.modes3], self.weights3)
      	out_ft[:, :, -self.modes1:, -self.modes2:, :self.modes3] = \
      	self.compl_mul3d(x_ft[:, :, -self.modes1:, -self.modes2:, :self.modes3], self.weights4)
      	#Return to physical space
      	x = torch.fft.irfftn(out_ft, s=(x.size(-3), x.size(-2), x.size(-1)))

  • How does the advantage of operator learning functions in this lab?
    • Does it not need to train again with different intial condition or boundrary condition.
    • Perhaps. As the Operator is learnt. So one of the three condtion has been solved, with the other two to be the intial condition and boundrary condition.
• Try with FNO for basic set up of the experiment.
• The FNO model for solving entire family of PDEs.
  • FNO project
  • FNO GitHub
  • FNO introduction
  • FNO data
  • FNO models
    • Here are the pre-trained models. It can be evaluated using eval.py or super_resolution.py
• By now there is no need to focus on the Galerkin Transformer, as it only provides a better precision on results.
• We will focus more on the FNO in the next stages.
• The FNO is more focused on solving a family of PDEs without try to train the neural network from the scratch, such as when the initial condition or the boundary condition change.