GSOC'25 Project: Implement a Wishlist Model - Fourier Neural Operator

#1

Hi!

I’m Shrisharanyan Vasu, a third-year undergraduate at Amrita Vishwa Vidyapeetham. I’m excited to contribute to DeepChem as part of the GSOC program for the year 2025. This summer, I will be implementing a Fourier Neural Operator class for efficiently solving partial differential equations by leveraging neural networks in the frequency domain.
As part of my contributions, I will be sharing weekly updates on the project’s progress, challenges, and milestones. I hope to engage with the community for feedback, suggestions, and collaboration to make this project a success.

Looking forward to working together and learning from all of you!

Regards,
Shrisharanyan Vasu

#2

Coding Period has begun!
Week 1 (Jun 2 - Jun 9)

  • Tested and got the SpectralConv layer merged into the codebase sometime back. Link to PR #4421
  • Read the FNO research paper and understood the basic implementation strategy.
  • Looked up and tried understanding the original FNO implementation: Link
  • Implemented a basic FNO model. Link to Draft PR #4452
  • Learnt how meshgrids help with preserving local positional information during training and tried searching for some existing implementations for FNO or some other fourier-based model which has meshgrids.

Link to slide decks: https://docs.google.com/presentation/d/1OLboTzxzKpvV6n42-kELSdOcWIooe0OuLAk5Ol74cow/edit?usp=sharing

#3

Week 2 (Jun 9 - Jun 16)

  • Added unit tests for the basic FNO model. Got PR (from previous week) reviewed by mentor and made the necessary changes. Link to commit
  • Working on implementing positional embedding using meshgrids to compensate for the lack of inherent spatial locality in spectral convolutions. Took inspiration from the original FNO implementation.

Link to slide decks: https://docs.google.com/presentation/d/1OLboTzxzKpvV6n42-kELSdOcWIooe0OuLAk5Ol74cow/edit?usp=sharing

#4

Week 3 (Jun 16 - Jun 20)

  • Split the previous PR into 2 PRs: one for the FNOBlock and another for the FNOModel. Added tests for the FNOBlock too.
  • Implemented meshgrids and opened a PR
  • Totally 3 PRs opened last week and this week:
    1. FNOBlock: PR #4459
    2. FNOModel: PR #4452
    3. Meshgrids: PR #4460

Link to slide decks: https://docs.google.com/presentation/d/1OLboTzxzKpvV6n42-kELSdOcWIooe0OuLAk5Ol74cow/edit?usp=sharing

#5

Week 4 (Jun 23 - Jun 27)

  • Removed a redundant test case in the FNOBlock PR.

  • Resolved comments from one of the mentors and opened new PRs to fix git commit leakage.

    1. FNOBlock: PR #4466
    2. FNOModel: PR #4467
    3. FNOModel with Meshgrids: PR #4468

  • Tested FNOModel on synthetic dataset (Link to Colab), created by generating samples using the equation: u += alpha * dt * (u_xx + u_yy)

Link to slide decks: https://docs.google.com/presentation/d/1OLboTzxzKpvV6n42-kELSdOcWIooe0OuLAk5Ol74cow/edit?usp=sharing

#6

Week 5 (Jun 30 - Jul 4)

  • Resolved documentation issues in the FNOBlock PR based on feedback from a mentor.
  • Looked at the different datasets available in PDEBench.
  • Wasn’t able to work much as I had semester final exams. Will catch up over the weekend and next week.
  • PRs open:
    – FNOModel: PR #4452
    – Meshgrids: PR #4460

To-do:

  1. Download Burgers Equation 1D and Navier Stokes 2D equation and test the FNOModel.
  2. Also test the original FNO implementation using the dataset and compare performance.

Link to slide decks: https://docs.google.com/presentation/d/1OLboTzxzKpvV6n42-kELSdOcWIooe0OuLAk5Ol74cow/edit?usp=sharing

#7

Week 6 (Jul 7 - Jul 11)

  • Got the FNOBlock PR merged in.
  • Debugged the differences between original implementation and my implementation. Made the following changes to my code:
    1. Added FFTShift before slicing the modes to obtain modes with frequencies close to zero, in SpectralConv.
    2. Included mode selection logic for modes centered around zero-frequency.
    3. Changed from ReLU to GeLU.
    4. Added normalization to inputs as part of data processing.
    5. Made the lifting and projection layers more robust by increasing the depth.

For results, click here for Colab Notebook

Link to slide decks: https://docs.google.com/presentation/d/1OLboTzxzKpvV6n42-kELSdOcWIooe0OuLAk5Ol74cow/edit?usp=sharing

#8

Week 7 (Jul 14 - Jul 18)

  • Made a patch PR for SpectralConv layer that fixes weight initialization.
  • Finished implementing a working prototype in colab.
  • Made a PR for the GaussianNormalizer class and FNO torch.nn.Module class.
  • Open PRs:
    1. SpectralConv Patch PR
    2. GaussianNormalizer and nn.Module FNO PR
    3. TorchModel FNOModel PR

Link to slide decks: https://docs.google.com/presentation/d/1OLboTzxzKpvV6n42-kELSdOcWIooe0OuLAk5Ol74cow/edit?usp=sharing

#9

Week 8 (Jul 21 - Jul 25)

  • Merged in the SpectralConv patch PR.
  • Addressed the comments left by mentor on the FNO nn.Module class.
  • Working on a tutorial for using the FNOModel TorchModel class.
  • Open PRs:
    1. GaussianNormalizer and nn.Module FNO PR
    2. TorchModel FNOModel PR

Click here for Colab Notebook (tutorial)

#10

Week 9 (Jul 28 - Aug 1)

  • Got nn.Module PR re-reviewed by my mentor.
  • Went through DeepChem’s NormalizationTransformer.
  • Working on a tutorial for using the FNOModel TorchModel class.
  • Open PRs:
    1. GaussianNormalizer and nn.Module FNO PR
    2. TorchModel FNOModel PR

Click here for Colab Notebook (tutorial)

Link to slide decks: https://docs.google.com/presentation/d/1OLboTzxzKpvV6n42-kELSdOcWIooe0OuLAk5Ol74cow/edit?usp=sharing

#11

Week 10 (Aug 4 - Aug 8)

  • Got LGTM from a mentor on nn.Module PR.
  • Fixed a bug which was causing the loss function to always be MSELoss in TorchModel PR.
  • Continued making progress on the tutorial.
  • Open PRs:
    1. nn.Module FNO PR
    2. TorchModel FNOModel PR

Click here for Colab Notebook containing tutorial

Link to slide decks: https://docs.google.com/presentation/d/1OLboTzxzKpvV6n42-kELSdOcWIooe0OuLAk5Ol74cow/edit?usp=sharing