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/src/NeuralOperators.jl b/src/NeuralOperators.jl index a4b1466..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 + 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 diff --git a/src/models/cno.jl b/src/models/cno.jl new file mode 100644 index 0000000..0329a22 --- /dev/null +++ b/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 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/models/cno_tests.jl b/test/models/cno_tests.jl new file mode 100644 index 0000000..a3f50c9 --- /dev/null +++ b/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 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" 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