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

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

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

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

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

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

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

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