feat: implement ConvolutionalNeuralOperator (CNO) (closes #122)

Project.toml

@@ -1,7 +1,7 @@
 name = "NeuralOperators"
 uuid = "ea5c82af-86e5-48da-8ee1-382d6ad7af4b"
 authors = ["Avik Pal <avikpal@mit.edu>"]
-version = "0.7.3"
+version = "0.8.0"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"

docs/Project.toml

@@ -29,7 +29,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"
 PythonCall = "0.9.23"

docs/src/api.md

@@ -9,6 +9,7 @@ and the Fourier transform used by spectral layers.
 NOMAD
 DeepONet
 FourierNeuralOperator
+ConvolutionalNeuralOperator
 ```
 
 ## Building blocks

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, 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
 

src/models/cno.jl

@@ -0,0 +1,187 @@
+"""
+    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 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.
+
+Resampling uses the FFT, so the spatial dimensions are treated as periodic.
+
+## 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 at the upsampled resolution.
+
+## Keyword Arguments
+
+  - `upsample_factor`: Integer upsampling factor of the activation operator. 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}
+    kernel = ntuple(Returns(3), N)
+    return CNOBlock(
+        Chain(
+            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},
+        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 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 conv → anti-aliased activation
+3. **Projection**: `Conv(1×…×1)` → anti-aliased activation → `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 by the anti-aliased activations. Default is
+    `gelu`.
+  - `upsample_factor`: Upsampling factor of the anti-aliased activations. 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),
+        CNOActivation(activation, upsample_factor),
+        Conv(ones_kernel, hidden_channels => out_channels),
+    )
+
+    return ConvolutionalNeuralOperator(Chain(; lifting, cno_blocks, projection))
+end

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"

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"

test/models/cno_tests.jl

@@ -0,0 +1,88 @@
+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
+
+    @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

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"

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