From 9931fed5fc228e6061c0c4f021bebc8b6d17bcd9 Mon Sep 17 00:00:00 2001 From: Jitendra Verma Date: Mon, 11 May 2026 10:05:36 +0530 Subject: [PATCH 1/4] feat: implement ConvolutionalNeuralOperator (CNO) (#122) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implements the Convolutional Neural Operator from: Raonic et al., "Convolutional Neural Operators for robust and accurate learning of PDEs", NeurIPS 2023. https://arxiv.org/abs/2302.01178 Architecture: - Lifting: Conv(1x...x1): in_channels → hidden_channels - CNO blocks: Upsample(:bilinear) → Conv(3x...x3, act, SamePad) → MeanPool - Projection: Conv(1x...x1, act) → Conv(1x...x1): → out_channels Each CNOBlock upsamples by upsample_factor, convolves at higher resolution, then downsamples via MeanPool — ensuring the operator converges to a continuous limit as spatial resolution increases (resolution-invariant). New types exported: - ConvolutionalNeuralOperator (the full model) - CNOBlock (the building block, composable) Closes #122 --- src/NeuralOperators.jl | 5 +- src/models/cno.jl | 151 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 155 insertions(+), 1 deletion(-) create mode 100644 src/models/cno.jl diff --git a/src/NeuralOperators.jl b/src/NeuralOperators.jl index a4b1466..ca13334 100644 --- a/src/NeuralOperators.jl +++ b/src/NeuralOperators.jl @@ -7,7 +7,7 @@ import FFTW using Random: Random, AbstractRNG using Lux: Lux, Chain, Dense, Conv, Parallel, NoOpLayer, WrappedFunction, Scale, - recursive_eltype + recursive_eltype, Upsample, MeanPool, SamePad using LuxCore: LuxCore, AbstractLuxLayer, AbstractLuxWrapperLayer using LuxLib: fast_activation!! using NNlib: batched_mul, gelu, pad_constant, sigmoid, sigmoid_fast, tanh_fast @@ -23,6 +23,7 @@ include("layers.jl") include("models/fno.jl") include("models/deeponet.jl") include("models/nomad.jl") +include("models/cno.jl") include("precompilation.jl") export FourierTransform @@ -32,6 +33,8 @@ export GridEmbedding, ComplexDecomposedLayer, SoftGating export FourierNeuralOperator export DeepONet export NOMAD +export ConvolutionalNeuralOperator + @public AbstractTransform, transform, truncate_modes, inverse diff --git a/src/models/cno.jl b/src/models/cno.jl new file mode 100644 index 0000000..df160f0 --- /dev/null +++ b/src/models/cno.jl @@ -0,0 +1,151 @@ +""" + CNOBlock( + in_channels::Integer, + out_channels::Integer, + modes::Dims{N}, + activation = gelu; + upsample_factor::Integer = 2, + ) where {N} + +A single Convolutional Neural Operator (CNO) block. + +Each block: +1. **Upsamples** the input by `upsample_factor` using bilinear/trilinear interpolation. +2. Applies a **3×(…×3) convolution** in the higher-resolution space with `SamePad`. +3. Applies the **activation function** pointwise. +4. **Downsamples** back to the original resolution via average pooling. + +This design ensures the discrete operator converges to a continuous limit as resolution +increases, unlike standard CNNs which only approximate finite-dimensional maps. + +## Arguments + + - `in_channels`: Number of input channels. + - `out_channels`: Number of output channels. + - `modes`: Spatial dimensions tuple (length = data dimensionality). Only the length is + used to set the kernel dimensionality. + - `activation`: Pointwise activation applied after convolution. + +## Keyword Arguments + + - `upsample_factor`: Integer upsampling factor. Default is `2`. + +## References + +[1] Raonic et al., "Convolutional Neural Operators for robust and accurate learning of +PDEs," NeurIPS 2023. https://arxiv.org/abs/2302.01178 +""" +@concrete struct CNOBlock <: AbstractLuxWrapperLayer{:model} + model +end + +function CNOBlock( + in_channels::Integer, + out_channels::Integer, + modes::Dims{N}, + activation = gelu; + upsample_factor::Integer = 2, + ) where {N} + spatial_dims = ntuple(identity, N) + kernel = ntuple(Returns(3), N) + pool_window = ntuple(Returns(upsample_factor), N) + + return CNOBlock( + Chain( + # 1. Upsample to higher resolution + Upsample(:bilinear; scale = upsample_factor), + # 2. Convolution at high resolution (preserves spatial size via SamePad) + Conv(kernel, in_channels => out_channels, activation; pad = SamePad()), + # 3. Downsample back via average pooling (paper Section 3.1) + MeanPool(pool_window), + ), + ) +end + +""" + ConvolutionalNeuralOperator( + modes::Dims{N}, + in_channels::Integer, + out_channels::Integer, + hidden_channels::Integer; + num_layers::Integer = 4, + activation = gelu, + upsample_factor::Integer = 2, + ) where {N} + +Convolutional Neural Operator (CNO) for learning PDE solution operators. + +CNO applies a sequence of resolution-preserving continuous convolutional blocks. Each +block upsamples the input, applies a convolution in the higher-resolution space, and +downsamples back. This design is proven to converge to a well-defined continuous operator +as resolution increases, making CNO resolution-invariant by construction. + +**Architecture**: +1. **Lifting** `Conv(1×…×1)`: maps `in_channels → hidden_channels` +2. **CNO blocks** × `num_layers`: each is upsample → conv → activation → avgpool +3. **Projection**: `Conv(1×…×1, act)` → `Conv(1×…×1)` maps to `out_channels` + +## Arguments + + - `modes`: Spatial size tuple (length `d` for d-dimensional data). Only its length + matters — kept consistent with the FNO API. + - `in_channels`: Number of input channels. + - `out_channels`: Number of output channels. + - `hidden_channels`: Number of channels inside the CNO blocks. + +## Keyword Arguments + + - `num_layers`: Number of `CNOBlock` layers. Default is `4`. + - `activation`: Activation function used inside each block. Default is `gelu`. + - `upsample_factor`: Spatial upsampling factor inside each block. Default is `2`. + +## References + +[1] Raonic et al., "Convolutional Neural Operators for robust and accurate learning of +PDEs," NeurIPS 2023. https://arxiv.org/abs/2302.01178 + +## Example + +```jldoctest +julia> cno = ConvolutionalNeuralOperator((16,), 1, 1, 32; num_layers=3); + +julia> ps, st = Lux.setup(Xoshiro(), cno); + +julia> u = rand(Float32, 64, 1, 5); + +julia> size(first(cno(u, ps, st))) +(64, 1, 5) +``` +""" +@concrete struct ConvolutionalNeuralOperator <: AbstractLuxWrapperLayer{:model} + model <: AbstractLuxLayer +end + +function ConvolutionalNeuralOperator( + modes::Dims{N}, + in_channels::Integer, + out_channels::Integer, + hidden_channels::Integer; + num_layers::Integer = 4, + activation = gelu, + upsample_factor::Integer = 2, + ) where {N} + ones_kernel = ntuple(Returns(1), N) + + lifting = Conv(ones_kernel, in_channels => hidden_channels) + + cno_blocks = Chain( + [ + CNOBlock( + hidden_channels, hidden_channels, modes, activation; upsample_factor, + ) for _ in 1:num_layers + ]..., + ) + + projection = Chain( + Conv(ones_kernel, hidden_channels => hidden_channels, activation), + Conv(ones_kernel, hidden_channels => out_channels), + ) + + return ConvolutionalNeuralOperator(Chain(; lifting, cno_blocks, projection)) +end From e3776787d117c12330b886c4733a4630e25d5721 Mon Sep 17 00:00:00 2001 From: Jitendra Verma Date: Mon, 18 May 2026 21:22:15 +0530 Subject: [PATCH 2/4] test: add CI-compatible Reactant gradient tests for ConvolutionalNeuralOperator --- test/models/cno_tests.jl | 60 ++++++++++++++++++++++++++++++++++++++++ test/runtests.jl | 3 ++ 2 files changed, 63 insertions(+) create mode 100644 test/models/cno_tests.jl diff --git a/test/models/cno_tests.jl b/test/models/cno_tests.jl new file mode 100644 index 0000000..e32b0c0 --- /dev/null +++ b/test/models/cno_tests.jl @@ -0,0 +1,60 @@ +using NeuralOperators, Test + +include("../shared_testsetup.jl") + +@testset "Convolutional Neural Operator" begin + rng = StableRNG(12345) + + setups = [ + ( + modes = (4,), + in_channels = 1, + out_channels = 1, + hidden_channels = 8, + num_layers = 2, + x_size = (16, 1, 4), + y_size = (16, 1, 4), + ), + ( + modes = (4, 4), + in_channels = 2, + out_channels = 1, + hidden_channels = 8, + num_layers = 2, + x_size = (16, 16, 2, 4), + y_size = (16, 16, 1, 4), + ), + ] + + xdev = reactant_device(; force = true) + + @testset "$(length(setup.modes))D" for setup in setups + cno = ConvolutionalNeuralOperator( + setup.modes, setup.in_channels, setup.out_channels, setup.hidden_channels; + num_layers = setup.num_layers, + ) + display(cno) + ps, st = Lux.setup(rng, cno) + + x = rand(rng, Float32, setup.x_size...) + y = rand(rng, Float32, setup.y_size...) + + @test size(first(cno(x, ps, st))) == setup.y_size + + ps_ra, st_ra = (ps, st) |> xdev + x_ra, y_ra = (x, y) |> xdev + + res = first(cno(x, ps, st)) + res_ra, _ = @jit cno(x_ra, ps_ra, st_ra) + @test res_ra ≈ res atol = 1.0f-2 rtol = 1.0f-2 + + @testset "check gradients" begin + ∂x_fd, ∂ps_fd = ∇sumabs2_finite_difference(cno, x, ps, st) + ∂x_ra, ∂ps_ra = ∇sumabs2_reactant(cno, x_ra, ps_ra, st_ra) + ∂x_ra, ∂ps_ra = (∂x_ra, ∂ps_ra) |> cpu_device() + + @test ∂x_fd ≈ ∂x_ra atol = 1.0f-1 rtol = 1.0f-1 + @test check_approx(∂ps_fd, ∂ps_ra; atol = 1.0f-1, rtol = 1.0f-1) + end + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 81df6de..fe07876 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -20,6 +20,9 @@ withenv( @time @safetestset "NOMAD" begin include(joinpath(@__DIR__, "models", "nomad_tests.jl")) end + @time @safetestset "Convolutional Neural Operator" begin + include(joinpath(@__DIR__, "models", "cno_tests.jl")) + end @time @safetestset "SpectralConv" begin include(joinpath(@__DIR__, "layers", "spectral_conv_tests.jl")) end From 872e08b740fef1507b7ac96868d5299e9132f7b9 Mon Sep 17 00:00:00 2001 From: Jitendra Verma Date: Sun, 27 Sep 2026 00:44:10 +0530 Subject: [PATCH 3/4] Document ConvolutionalNeuralOperator in API and bump to 0.8.0 --- Project.toml | 2 +- docs/Project.toml | 2 +- docs/src/api.md | 1 + test/Project.toml | 2 +- test/gpu/Project.toml | 2 +- test/qa/Project.toml | 2 +- 6 files changed, 6 insertions(+), 5 deletions(-) diff --git a/Project.toml b/Project.toml index 87950cc..cbaadb3 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "NeuralOperators" uuid = "ea5c82af-86e5-48da-8ee1-382d6ad7af4b" authors = ["Avik Pal "] -version = "0.7.3" +version = "0.8.0" [deps] AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c" diff --git a/docs/Project.toml b/docs/Project.toml index 24eb43e..f6ab252 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -26,7 +26,7 @@ Git = "1.5.0" Lux = "1" MAT = "0.10.7, 0.11, 0.12" MLUtils = "0.4.4" -NeuralOperators = "0.7" +NeuralOperators = "0.7, 0.8" Optimisers = "0.4" Printf = "1.10, < 0.0.1, 1" Reactant = "0.2.130" diff --git a/docs/src/api.md b/docs/src/api.md index 0dc3db1..41e5cb8 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -9,6 +9,7 @@ and the Fourier transform used by spectral layers. NOMAD DeepONet FourierNeuralOperator +ConvolutionalNeuralOperator ``` ## Building blocks diff --git a/test/Project.toml b/test/Project.toml index f00eb60..772aec7 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -27,7 +27,7 @@ LuxCore = "1" LuxLib = "1.2" MLDataDevices = "1.17" NNlib = "0.9" -NeuralOperators = "0.6, 0.7" +NeuralOperators = "0.6, 0.7, 0.8" Optimisers = "0.4" Random = "1.10, < 0.0.1, 1" Reactant = "0.2.239" diff --git a/test/gpu/Project.toml b/test/gpu/Project.toml index d58fec0..82efe0e 100644 --- a/test/gpu/Project.toml +++ b/test/gpu/Project.toml @@ -32,7 +32,7 @@ LuxCore = "1" LuxLib = "1.2" MLDataDevices = "1.17" NNlib = "0.9" -NeuralOperators = "0.7" +NeuralOperators = "0.7, 0.8" Optimisers = "0.4" Random = "1.10" Reactant = "0.2.239" diff --git a/test/qa/Project.toml b/test/qa/Project.toml index 6aac4f9..4db7d55 100644 --- a/test/qa/Project.toml +++ b/test/qa/Project.toml @@ -16,7 +16,7 @@ Documenter = "1.5.0" FastTransforms = "0.17.1" JET = "0.9, 0.10, 0.11, 0.12" Lux = "1" -NeuralOperators = "0.7" +NeuralOperators = "0.7, 0.8" Random = "1.10" SafeTestsets = "0.1, 1" SciMLTesting = "2.4" From 9def10b6c4fda9849853c27f33dd482594e700d1 Mon Sep 17 00:00:00 2001 From: Jitendra Verma Date: Thu, 1 Oct 2026 01:17:25 +0530 Subject: [PATCH 4/4] Use anti-aliased activation in CNO blocks Replace bilinear upsampling and mean pooling with band-limited FFT resampling around the activation, following the paper's activation operator, and use the same filtered activation in the projection. Add an aliasing check and a spectral resampling test. --- src/NeuralOperators.jl | 2 +- src/models/cno.jl | 90 ++++++++++++++++++++++++++++------------ test/models/cno_tests.jl | 28 +++++++++++++ 3 files changed, 92 insertions(+), 28 deletions(-) diff --git a/src/NeuralOperators.jl b/src/NeuralOperators.jl index ca13334..632beed 100644 --- a/src/NeuralOperators.jl +++ b/src/NeuralOperators.jl @@ -7,7 +7,7 @@ import FFTW using Random: Random, AbstractRNG using Lux: Lux, Chain, Dense, Conv, Parallel, NoOpLayer, WrappedFunction, Scale, - recursive_eltype, Upsample, MeanPool, SamePad + recursive_eltype, SamePad using LuxCore: LuxCore, AbstractLuxLayer, AbstractLuxWrapperLayer using LuxLib: fast_activation!! using NNlib: batched_mul, gelu, pad_constant, sigmoid, sigmoid_fast, tanh_fast diff --git a/src/models/cno.jl b/src/models/cno.jl index df160f0..0329a22 100644 --- a/src/models/cno.jl +++ b/src/models/cno.jl @@ -9,14 +9,13 @@ A single Convolutional Neural Operator (CNO) block. -Each block: -1. **Upsamples** the input by `upsample_factor` using bilinear/trilinear interpolation. -2. Applies a **3×(…×3) convolution** in the higher-resolution space with `SamePad`. -3. Applies the **activation function** pointwise. -4. **Downsamples** back to the original resolution via average pooling. +Each block applies a `3×(…×3)` convolution followed by the CNO activation operator: the +signal is upsampled by `upsample_factor` with band-limited (sinc) interpolation, the +activation is applied pointwise at the higher resolution, and the result is low-pass +filtered and downsampled back to the input resolution. Filtering before downsampling keeps +the harmonics created by the activation from aliasing into lower frequencies. -This design ensures the discrete operator converges to a continuous limit as resolution -increases, unlike standard CNNs which only approximate finite-dimensional maps. +Resampling uses the FFT, so the spatial dimensions are treated as periodic. ## Arguments @@ -24,11 +23,11 @@ increases, unlike standard CNNs which only approximate finite-dimensional maps. - `out_channels`: Number of output channels. - `modes`: Spatial dimensions tuple (length = data dimensionality). Only the length is used to set the kernel dimensionality. - - `activation`: Pointwise activation applied after convolution. + - `activation`: Pointwise activation applied at the upsampled resolution. ## Keyword Arguments - - `upsample_factor`: Integer upsampling factor. Default is `2`. + - `upsample_factor`: Integer upsampling factor of the activation operator. Default is `2`. ## References @@ -46,22 +45,55 @@ function CNOBlock( activation = gelu; upsample_factor::Integer = 2, ) where {N} - spatial_dims = ntuple(identity, N) kernel = ntuple(Returns(3), N) - pool_window = ntuple(Returns(upsample_factor), N) - return CNOBlock( Chain( - # 1. Upsample to higher resolution - Upsample(:bilinear; scale = upsample_factor), - # 2. Convolution at high resolution (preserves spatial size via SamePad) - Conv(kernel, in_channels => out_channels, activation; pad = SamePad()), - # 3. Downsample back via average pooling (paper Section 3.1) - MeanPool(pool_window), + Conv(kernel, in_channels => out_channels; pad = SamePad()), + CNOActivation(activation, upsample_factor), ), ) end +@concrete struct CNOActivation <: AbstractLuxLayer + activation + upsample_factor::Int +end + +function (act::CNOActivation)(x::AbstractArray{T, M}, _, st::NamedTuple) where {T, M} + sz = size(x)[1:(M - 2)] + y = act.activation.(spectral_resample(x, sz .* act.upsample_factor)) + return spectral_resample(y, sz), st +end + +# Band-limited resampling of the spatial dimensions of `x` to size `sz`. Frequencies at or +# above the lower of the two Nyquist limits are dropped. +function spectral_resample(x::AbstractArray{T}, sz::Dims{N}) where {T, N} + in_sz = size(x)[1:N] + in_sz == sz && return x + x_fft = resize_half_spectrum(rfft(x, 1:N), first(in_sz), first(sz)) + for d in 2:N + x_fft = resize_full_spectrum(x_fft, d, sz[d]) + end + y = irfft(x_fft, first(sz), 1:N) + return y .* T(length(y) / length(x)) +end + +function select_range(x::AbstractArray, d::Int, r::AbstractUnitRange) + return x[ntuple(i -> i == d ? r : Colon(), ndims(x))...] +end + +function resize_half_spectrum(x_fft::AbstractArray, n::Int, m::Int) + k = (min(n, m) - 1) ÷ 2 + return pad_constant(select_range(x_fft, 1, 1:(k + 1)), (0, m ÷ 2 - k), false; dims = 1) +end + +function resize_full_spectrum(x_fft::AbstractArray, d::Int, m::Int) + n = size(x_fft, d) + k = (min(n, m) - 1) ÷ 2 + pos = pad_constant(select_range(x_fft, d, 1:(k + 1)), (0, m - 2k - 1), false; dims = d) + return cat(pos, select_range(x_fft, d, (n - k + 1):n); dims = d) +end + """ ConvolutionalNeuralOperator( modes::Dims{N}, @@ -75,15 +107,17 @@ end Convolutional Neural Operator (CNO) for learning PDE solution operators. -CNO applies a sequence of resolution-preserving continuous convolutional blocks. Each -block upsamples the input, applies a convolution in the higher-resolution space, and -downsamples back. This design is proven to converge to a well-defined continuous operator -as resolution increases, making CNO resolution-invariant by construction. +CNO applies a sequence of resolution-preserving blocks. Each block is a convolution +followed by an anti-aliased activation: the signal is upsampled with band-limited (sinc) +interpolation, the activation is applied at the higher resolution, and the result is +low-pass filtered and downsampled back. Resampling uses the FFT, so the spatial dimensions +are treated as periodic. **Architecture**: 1. **Lifting** `Conv(1×…×1)`: maps `in_channels → hidden_channels` -2. **CNO blocks** × `num_layers`: each is upsample → conv → activation → avgpool -3. **Projection**: `Conv(1×…×1, act)` → `Conv(1×…×1)` maps to `out_channels` +2. **CNO blocks** × `num_layers`: each is conv → anti-aliased activation +3. **Projection**: `Conv(1×…×1)` → anti-aliased activation → `Conv(1×…×1)` maps to + `out_channels` ## Arguments @@ -96,8 +130,9 @@ as resolution increases, making CNO resolution-invariant by construction. ## Keyword Arguments - `num_layers`: Number of `CNOBlock` layers. Default is `4`. - - `activation`: Activation function used inside each block. Default is `gelu`. - - `upsample_factor`: Spatial upsampling factor inside each block. Default is `2`. + - `activation`: Activation function used by the anti-aliased activations. Default is + `gelu`. + - `upsample_factor`: Upsampling factor of the anti-aliased activations. Default is `2`. ## References @@ -143,7 +178,8 @@ function ConvolutionalNeuralOperator( ) projection = Chain( - Conv(ones_kernel, hidden_channels => hidden_channels, activation), + Conv(ones_kernel, hidden_channels => hidden_channels), + CNOActivation(activation, upsample_factor), Conv(ones_kernel, hidden_channels => out_channels), ) diff --git a/test/models/cno_tests.jl b/test/models/cno_tests.jl index e32b0c0..a3f50c9 100644 --- a/test/models/cno_tests.jl +++ b/test/models/cno_tests.jl @@ -57,4 +57,32 @@ include("../shared_testsetup.jl") @test check_approx(∂ps_fd, ∂ps_ra; atol = 1.0f-1, rtol = 1.0f-1) end end + + @testset "activation does not alias" begin + block = NeuralOperators.CNOBlock(1, 1, (4,), abs2) + ps, st = Lux.setup(rng, block) + ps.layer_1.weight .= 0 + ps.layer_1.weight[2, 1, 1] = 1 + ps.layer_1.bias .= 0 + + # cos² of frequency 12 is DC plus frequency 24, which must be dropped, not folded + # back to frequency 8, where the grid cannot represent it + for n in (32, 64, 128) + x = reshape(Float32[cos(2π * 12 * j / n) for j in 0:(n - 1)], n, 1, 1) + y = vec(first(block(x, ps, st))) + @test sum(y) / n ≈ 0.5 atol = 1.0f-5 + @test 2abs(sum(y .* cis.(-2π * 8 * (0:(n - 1)) / n))) / n < 1.0f-5 + end + end + + @testset "spectral resampling" begin + f(x, y) = cos(2π * (3x + 5y)) + sin(2π * (-4x + 2y)) + samples(n) = reshape([f(i / n[1], j / n[2]) for i in 0:(n[1] - 1), j in 0:(n[2] - 1)], n..., 1, 1) + + for (n, m) in (((16, 12), (32, 36)), ((15, 11), (30, 22)), ((16, 12), (13, 17))) + y = NeuralOperators.spectral_resample(samples(n), m) + @test y ≈ samples(m) + @test NeuralOperators.spectral_resample(y, n) ≈ samples(n) + end + end end