Some docs here would be helpful. Please refer to the contributors guide for guidance on how to format in-line docstrings

These layers belong in neuralop/layers

same comments re: docstring, and moving layers to the layers module

This is great. Let's think of a more descriptive name than MWT, perhaps MultiWaveletNO or MWNO?

We don't necessarily need these classes

Do we need to go through Sympy for this? Couldn't we directly return the discretized evaluations, instead of generating a symbolic version and then evaluating?

Yes, we use numpy instead of Sympy to make the calculation faster.

These are extremely helpful docstrings. I would suggest putting them down beneath the params heading so that the parameters section renders first in IDEs

Thank you! I have moved the place for those senteces.

Let's move the description back to the top of the docstring, I think it's what we have been doing more consistently across the library. Sorry for the back and forth

Maybe we can remove these since they don't seem to provide more functionality than what n_modes does?

These (alpha and n_dim) are not in the __init__ anymore. If that's intentional, we can remove them from the doctoring so users are not confused

Let's make sure the docstring reflects only the parameters that the user can specify in the __init__

Could we use the existing ChannelMLP class (link) ?

You can then create the lifting layer as in the init of FNO class

self.lifting = ChannelMLP(
            in_channels=lifting_in_channels,
            out_channels=self.hidden_channels,
            hidden_channels=self.lifting_channels,
            n_layers=2,
            n_dim=self.n_dim,
            non_linearity=non_linearity,
        )
        if self.complex_data:
            self.lifting = ComplexValued(self.lifting)

Could we use the existing ChannelMLP class (link) ?

You can then create the projection layer as in the init of FNO class

## Projection layer
        self.projection = ChannelMLP(
            in_channels=self.hidden_channels,
            out_channels=out_channels,
            hidden_channels=self.projection_channels,
            n_layers=2,
            n_dim=self.n_dim,
            non_linearity=non_linearity,
        )
        if self.complex_data:
            self.projection = ComplexValued(self.projection)

It seems that the current code only allows to use the same number of modes for every dimension. Is there a real constraint leading to this design? Otherwise it would be good to have the flexibility of choosing more or less modes along certain dimensions, since dynamics can have different scales along different dimensions and the resolution could also be different along different dimensions

Thank you for your question, but currently, the code only allows the same number of modes for every dimension. We will improve this feature later.

I don't think this function is used anywhere

The projection_channels=0 option does not seem to be implemented. Maybe we can remove that option? Or simply use the ChannelMLP already in the library

Thanks — we now use ChannelMLP for projection (same pattern as FNO), with projection_channels=0 mapped to hidden width 128 so the documented default is actually applied; forward permutes to channels-first for the MLP and back.

Maybe can remove mention of alpha since we will remove it

Might be good to remove alpha completely everywhere and just use n_modes to avoid having two variables for the same thing

Dropped redundant alpha / duplicate self.alpha; the public API is n_modes only (with optional int shorthand for 1D). MWNOBlock now takes n_modes as well; internal kernels still use the first entry as the mode cutoff, documented in Notes.

With the current implementation, n_modes cannot really be a tuple. It seems the code only takes n_modes[0] and then uses it everywhere.

Either we allow for the number of modes to be different along different dimensions, or otherwise we update the docstring and an assertion at the beginning of the model definition to make it clear n_modes but be a single integer

Good point. The public API now requires n_modes to be one int plus explicit n_dim, with validation in init, and the docs no longer imply a tuple. Anisotropic mode counts would need deeper changes in SparseKernelFT; we can leave that for a later PR if desired.

Definitely need to add an assert and throw an error if one of the spatial resolutions is not a power of 2

Definitely need to add an assert and throw an error if one of the spatial resolutions is not a power of 2

Added _validate_spatial_resolution at the start of forward: we require power-of-two sizes on the wavelet axes (1D: n_points; 2D: H,W; 3D: H,W only) and raise ValueError with an explicit message otherwise, plus a small unit test for the 2D case.

phi does not take these lb and ub arguments anymore. Please update the names and make sure everything runs smoothly

phi_ uses lower_bound / upper_bound, not lb / ub. Updated the Chebyshev partial(...) calls accordingly so the wavelet callables match the current signature.

It does not seem like the Chebyshev version runs correctly. With the below code, it seems that psi1 and psi2 are just zeros. Make sure to run the code with the Chebyshev version and make sure it works properly

You were right: Chebyshev psi1/psi2 were zero polynomials. We now convert the Chebyshev scaling rows to power basis (so the same Gram–Schmidt + monomial integrals as Legendre are valid), reuse that Gram–Schmidt for Chebyshev wavelet coefficients, and build psi1/psi2 with phi_ and lower_bound/upper_bound on [0,0.5] / [0.5,1]. Added tests to assert non-zero wavelets and a finite Chebyshev filter bank.

These could be overlapping and problematic if alpha is too large. This is why I would recommend using the SpectralConv class directly, or starting from there

Thanks — we were indeed using m = min(α, nx//2+1) on the full FFT axis, so [:m] and [-m:] could overlap when 2m > nx. We now cap symmetric blocks at nx//2 per side (and treat nx≤1 as DC-only with no negative slice, since [-0:] aliases the whole axis). The last rfft dimension still uses n//2+1. We added a short note pointing to SpectralConv’s handling of FFT redundancy; wiring SparseKernelFT to SpectralConv directly would be a larger API/weight-shape change if we want that refactor later. This would require a complete structural modification.

Would be good to add an assert that L needs to be smaller than num_scales

Added a ValueError in MWNO.forward when L >= num_scales (with num_scales = floor(log2(spatial_size)) on the wavelet axis), documented the constraint on L for both MWNO and MWNOBlock, and a small unit test.

Need to test more cases, and definitely need to test chebyshev. If it becomes a number of configs that is too large, we can split it in 2-3 test functions

Thanks — I split the parametrized smoke test into test_mwno_legendre and test_mwno_chebyshev with the same grid (n_dim/n_modes, c, n_layers ∈ {1,3}, L ∈ {0,1,2}, k=4) and factored the forward/backward + grad checks into a small helper so we don’t duplicate code. If 72 parametrized runs is too heavy for CI we can trim (e.g. drop L=2 or one of n_layers).