[WIP] Integral Kernel Operator

Project.toml

@@ -7,7 +7,9 @@ version = "0.1.0"
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
+Cubature = "667455a9-e2ce-5579-9412-b964f529a492"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
+Integrals = "de52edbc-65ea-441a-8357-d3a637375a31"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"

src/LuxNeuralOperators.jl

@@ -12,6 +12,7 @@ using PrecompileTools: @recompile_invalidations
     using NNlib: NNlib, ⊠
     using Random: Random, AbstractRNG
     using Reexport: @reexport
+    using Integrals
 end
 
 const CRC = ChainRulesCore
@@ -26,10 +27,12 @@ include("layers.jl")
 
 include("fno.jl")
 include("deeponet.jl")
+include("iko.jl")
 
 export FourierTransform
 export SpectralConv, OperatorConv, SpectralKernel, OperatorKernel
 export FourierNeuralOperator
 export DeepONet
+export IntegralKernel, IntegralKernelOperator
 
 end

src/iko.jl

@@ -0,0 +1,82 @@
+
+function IntegralKernel(W::Tuple, κ::Tuple; # bias,
+        W_activation=identity, kernel_activation=identity,
+        activation=identity, alg=Integrals.HCubatureJL(), kwargs...)
+    W_ = Chain([Dense(W[i] => W[i + 1], W_activation) for i in 1:(length(W) - 1)]...)
+    κ_ = Chain([Dense(κ[i] => κ[i + 1], kernel_activation) for i in 1:(length(κ) - 1)]...)
+
+    IntegralKernel(W_, κ_; activation=activation, alg=alg, kwargs...)
+end
+
+"""
+    IntegralKernel(W::L1, κ::L2, domain; # bias,
+            activation=identity,
+            alg=Integrals.HCubatureJL(), kwargs...) where {L1, L2}
+
+returns the Integral kernel evaluated the given data point: ``σ(W_t + \\mathcal{K}_t)(v_t)(x)``
+where ``W_t`` is a linear mapping and
+
+```math
+(\\mathcal{K}_t(v_t))(x) = \\int_{D_t} \\kappa^{(t)}(x,y) v_t(y) dy \\quad \\forall x \\in D_t
+```
+
+## Arguments
+
+  - `W` : network for linear mapping
+  - `κ` : network to evaluate the integral kernel
+  - `domain` : domain of integration to perform integration of `κ`
+
+## Keyword arguments
+
+  - `activation` : activation function to be applied at the end σ # bias,
+  - `alg` : `Integrals.jl` algorithm to compute the integral
+  - `kwargs` : Additional arguments to be splatted into `Integrals.solve(...)`
+"""
+function IntegralKernel(W::L1, κ::L2; # bias,
+        activation=identity,
+        alg=Integrals.HCubatureJL(), kwargs...) where {L1, L2}
+
+    # name
+
+    return @compact(; W, κ, activation, alg, kwargs,
+        dispatch=:IntegralKernel) do (x, domain)
+        W_ = W(x)
+        f(u, p) = κ(vcat(u, x))
+        prototype = zero(x)
+        prob = IntegralProblem(IntegralFunction(f, prototype), domain)
+        sol = solve(prob, alg; kwargs...)
+        # print("wewe \n")
+        @return W_ #broadcast(activation, W_)
+    end
+end
+
+"""
+    function IntegralKernelOperator(
+        lifting::L1, kernels::Vector{L2},
+        projection::L3, domain) where {L1, L2 <: CompactLuxLayer{:IntegralKernel}, L3}
+
+returns the continuous variant of Neural Operator
+
+## Arguments
+
+  - `lifting`: lifting layer
+  - `kernels`: Vector of `IntegralKernel` to applied in chain after lifting
+  - `projection`: projection layer
+"""
+function IntegralKernelOperator(lifting::L1, kernels::Vector{L2}, projection::L3,
+        domain) where {L1, L2 <: CompactLuxLayer{:IntegralKernel}, L3}
+    return @compact(; lifting, kernels, projection, domain,
+        dispatch=:IntegralKernelOperator) do x
+        v = lifting(x)
+        D = sort(lifting(domain), dims = 2)
+
+        for kernel in kernels
+            v = kernel((v, D)) # kernel evaluation
+
+            D = sort(kernel((D, D)), dims = 2) # update domain of integration for next kernel
+        end
+
+        v = projection(v)
+        @return v
+    end
+end

test/integral_kernel_tests.jl

@@ -0,0 +1,44 @@
+@testitem "IntegralKernelOperator" setup=[SharedTestSetup] begin
+    @testset "BACKEND: $(mode)" for (mode, aType, dev, ongpu) in MODES
+        rng = StableRNG(12345)
+
+        @testset "Kernel test" begin
+            u = rand(Float64, 1, 16) |> aType
+            domain_ = reshape([0.0, 1.0], 1, 2) |> aType
+
+            integral_kernel_ = IntegralKernel(
+                Chain(Dense(1 => 16), Dense(16 => 16), Dense(16 => 1)),
+                Chain(Dense(2 => 16), Dense(16 => 16), Dense(16 => 1)))
+
+            ps, st = Lux.setup(rng, integral_kernel_) |> dev
+
+            # @inferred integral_kernel_((u, domain_), ps, st)
+            # @jet integral_kernel_((u, domain_), ps, st)
+
+            pred = first(integral_kernel_((u, domain_), ps, st))
+            @test size(pred) == size(u)
+        end
+
+        @testset "Operator test" begin
+            u = rand(Float64, 1, 16) |> aType
+            domain_ = reshape([0.0, 1.0], 1, 2) |> aType
+
+            kernels = [IntegralKernel(
+                           Chain(Dense(1 => 16), Dense(16 => 16), Dense(16 => 1)),
+                           Chain(Dense(2 => 16), Dense(16 => 16), Dense(16 => 1)))
+                       for _ in 1:3]
+
+            model = IntegralKernelOperator(
+                Chain(Dense(1 => 16), Dense(16 => 16), Dense(16 => 1)), kernels,
+                Chain(Dense(1 => 16), Dense(16 => 16), Dense(16 => 1)), domain_)
+
+            ps, st = Lux.setup(rng, model) |> dev
+
+            # @inferred model(u, ps, st)
+            # @jet model(u, ps, st)
+
+            pred = first(model(u, ps, st))
+            @test size(pred) == size(u)
+        end
+    end
+end

test/shared_testsetup.jl

@@ -21,8 +21,8 @@ end
 const MODES = begin
     modes = []
     cpu_testing() && push!(modes, ("CPU", Array, LuxCPUDevice(), false))
-    cuda_testing() && push!(modes, ("CUDA", CuArray, LuxCUDADevice(), true))
-    amdgpu_testing() && push!(modes, ("AMDGPU", ROCArray, LuxAMDGPUDevice(), true))
+    # cuda_testing() && push!(modes, ("CUDA", CuArray, LuxCUDADevice(), true))
+    # amdgpu_testing() && push!(modes, ("AMDGPU", ROCArray, LuxAMDGPUDevice(), true))
     modes
 end