Skip to content

Add wavelet convolution layer#639

Open
sinapordanesh wants to merge 11 commits intoneuraloperator:mainfrom
sinapordanesh:feat/layers-wavelet-conv
Open

Add wavelet convolution layer#639
sinapordanesh wants to merge 11 commits intoneuraloperator:mainfrom
sinapordanesh:feat/layers-wavelet-conv

Conversation

@sinapordanesh
Copy link
Copy Markdown

@sinapordanesh sinapordanesh commented Aug 14, 2025

Wavelet Neural Operator (WNO)

Resource Link
Paper https://arxiv.org/pdf/2205.02191
Original implementation https://github.com/TapasTripura/WNO
Citation @Article{tripura2023wavelet,
title={Wavelet Neural Operator for solving parametric partial differential equations in computational mechanics problems},
author={Tripura, Tapas and Chakraborty, Souvik},
journal={Computer Methods in Applied Mechanics and Engineering},
volume={404},
pages={115783},
year={2023},
publisher={Elsevier}}

Spectral convolution class and its helper functions for operator learning in wavelet space. Main added class is "SpectralConvWavelet" in neuralop/layers/spectral_convolution_laplace.py, which handles 1-3 dimensions in an abstract way.

The code mostly follows the paper code, just reformed in an abstract way to handle all 3 dimensions in a single class, via helper functions for building einsums sentences, some other common features through different paper classes.
A set of test cases is included in this PR under neuralop/layers/tests/test_spectral_convolution_laplace.py to test the basic features of the forward method of the spectral convolution.

Comparison between paper code training results and the PR code results.

The PR code tested vs. regenerated results via code provided at the paper's repository at https://github.com/TapasTripura/WNO. We run all scenarios for which data was provided, as well as those runnable on our computational capacity.

Mean Testing Error

Run System Paper code results (%) PR code results (%)
1D_Advection 0.842 0.836
1D_Advection_Time 0.946 0.946
1D_Burgers 2.043 2.043
2D_Darcy_dwt 2.350 2.297
2D_Darcy_notch_dwt 2.190 2.050
2D_NS_dwt 4.625 4.0270

For testing the integrity of the PR spectral convolution class, we just exchanged the spectral convolution objects lists self.conv from the original class WaveConv<n>d to the PR class SpectralConvWavelet. The rest of the training scripts remained the same as the original.
All the trainng data can be found in these links:
Link1, Link2, Link3

Detailed results and running scripts are in the following zip file.
scipts_and_results.zip

@sinapordanesh
Copy link
Copy Markdown
Author

sinapordanesh commented Aug 14, 2025

Besides the code and test cases, this wavelet-based operation relies on two extra dependencies:

  1. "Pytorch Wavelets" (see https://pytorch-wavelets.readthedocs.io/en/latest/readme.html) (for 1&2D)
    $ git clone https://github.com/fbcotter/pytorch_wavelets
    $ cd pytorch_wavelets
    $ pip install .
  1. "Pytorch Wavelet Toolbox" (see https://github.com/v0lta/PyTorch-Wavelet-Toolbox) (for 3D)
    $ pip install ptwt

@neuraloperator neuraloperator deleted a comment from sinapordanesh Aug 14, 2025
@JeanKossaifi
Copy link
Copy Markdown
Member

Thank you @sinapordanesh - sorry for the delay in reviewing this. Seems there is a failing test.

@sinapordanesh
Copy link
Copy Markdown
Author

sinapordanesh commented Sep 9, 2025

Thank you @sinapordanesh - sorry for the delay in reviewing this. Seems there is a failing test.

Thank you, @JeanKossaifi, for reviewing this PR.
The failure is due to the dependency requirements. We need Wavelet Transform packages to execute the transformation on the signal before the neural operation. (We natively have FFT in the torch package, but no similar function for Wavelet Transform.)

I've mentioned required dependencies here:
#639 (comment)

Hopefully they are legitimate packages and can be included in your testing environment.

@JeanKossaifi
Copy link
Copy Markdown
Member

I agree, it seems useful to me.

I think we could make the import optional as we did with e.g. Spherical Harmonics? @dhpitt @vduruiss what do you think?

@dhpitt
Copy link
Copy Markdown
Member

dhpitt commented Sep 10, 2025 via email

- abstract funciton to form an einsum multiplication epxression, based on the dimension of the operation.
- branching 3 different spectral convolution for 1-3dim
- wavelt transform for 1&2d with DWT module
- wavelet transform for 3d with ptwt_wavedec3 module
- building einsum for 3d locally
- level adjust helper function for 1&2d
- einsum helper function to execute the 1&2d einsums
- 3d einsum seperated in the 3d branch at forward method.
- Un-debugged version of forward method. adopted directly from paper code -> Debug it!
@sinapordanesh sinapordanesh force-pushed the feat/layers-wavelet-conv branch 2 times, most recently from 3255beb to fea62f5 Compare October 14, 2025 21:49
@sinapordanesh
Copy link
Copy Markdown
Author

@JeanKossaifi, I fixed the dependency requirements and made it optional.
Let me know if there is anything else on both Wavelet and Laplace PRs.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants