refactor: rework the package to be Reactant first

.JuliaFormatter.toml

@@ -1,9 +1,3 @@
-style = "sciml"
-format_markdown = true
-whitespace_in_kwargs = false
-margin = 92
-indent = 4
-format_docstrings = true
-separate_kwargs_with_semicolon = true
-always_for_in = true
-annotate_untyped_fields_with_any = false
+style = "blue"
+pipe_to_function_call = false
+always_use_return = true

.buildkite/documentation.yml

@@ -2,7 +2,7 @@ steps:
   - label: ":julia: Documentation"
     plugins:
       - JuliaCI/julia#v1:
-          version: "1.10"
+          version: "1"
       - JuliaCI/julia-coverage#v1:
           codecov: true
     command: |

.buildkite/testing.yml

@@ -19,29 +19,29 @@ steps:
         julia:
           - "1"
 
-  - label: ":julia: Julia: {{matrix.julia}} + AMD GPU"
-    plugins:
-      - JuliaCI/julia#v1:
-          version: "{{matrix.julia}}"
-      - JuliaCI/julia-test#v1:
-          test_args: "--quickfail"
-      - JuliaCI/julia-coverage#v1:
-          codecov: true
-    env:
-      JULIA_AMDGPU_CORE_MUST_LOAD: "1"
-      JULIA_AMDGPU_HIP_MUST_LOAD: "1"
-      JULIA_AMDGPU_DISABLE_ARTIFACTS: "1"
-      BACKEND_GROUP: "AMDGPU"
-    agents:
-      queue: "juliagpu"
-      rocm: "*"
-      rocmgpu: "*"
-    if: build.message !~ /\[skip tests\]/
-    timeout_in_minutes: 60
-    matrix:
-      setup:
-        julia:
-          - "1"
+  # - label: ":julia: Julia: {{matrix.julia}} + AMD GPU"
+  #   plugins:
+  #     - JuliaCI/julia#v1:
+  #         version: "{{matrix.julia}}"
+  #     - JuliaCI/julia-test#v1:
+  #         test_args: "--quickfail"
+  #     - JuliaCI/julia-coverage#v1:
+  #         codecov: true
+  #   env:
+  #     JULIA_AMDGPU_CORE_MUST_LOAD: "1"
+  #     JULIA_AMDGPU_HIP_MUST_LOAD: "1"
+  #     JULIA_AMDGPU_DISABLE_ARTIFACTS: "1"
+  #     BACKEND_GROUP: "AMDGPU"
+  #   agents:
+  #     queue: "juliagpu"
+  #     rocm: "*"
+  #     rocmgpu: "*"
+  #   if: build.message !~ /\[skip tests\]/
+  #   timeout_in_minutes: 60
+  #   matrix:
+  #     setup:
+  #       julia:
+  #         - "1"
 
 env:
   RETESTITEMS_NWORKERS: 4

.github/workflows/CI.yml

@@ -29,10 +29,10 @@ jobs:
       matrix:
         version:
           - "1.10"
+          - "1"
         os:
           - ubuntu-latest
           - macos-latest
-          - windows-latest
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v2

.github/workflows/CompatHelper.yml

@@ -37,7 +37,7 @@ jobs:
       - name: "Run CompatHelper"
         run: |
           import CompatHelper
-          CompatHelper.main()
+          CompatHelper.main(; subdirs=["", "docs", "test"])
         shell: julia --color=yes {0}
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

Project.toml

@@ -1,33 +1,25 @@
 name = "NeuralOperators"
 uuid = "ea5c82af-86e5-48da-8ee1-382d6ad7af4b"
 authors = ["Avik Pal <avikpal@mit.edu>"]
-version = "0.5.3"
+version = "0.6.0"
 
 [deps]
-ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
-ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
-FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
 LuxLib = "82251201-b29d-42c6-8e01-566dec8acb11"
-MLDataDevices = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
-Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
 WeightInitializers = "d49dbf32-c5c2-4618-8acc-27bb2598ef2d"
 
 [compat]
-ArgCheck = "2.3"
-ChainRulesCore = "1.24"
+AbstractFFTs = "1.5.0"
 ConcreteStructs = "0.2.3"
-FFTW = "1.8"
-Lux = "1"
-LuxCore = "1"
-LuxLib = "1.2"
-MLDataDevices = "1.2.0"
-NNlib = "0.9.21"
+Lux = "1.13"
+LuxCore = "1.2"
+LuxLib = "1.8"
+NNlib = "0.9.30"
 Random = "1.10"
-Static = "1.1.1"
 WeightInitializers = "1"
 julia = "1.10"

docs/Project.toml

@@ -1,29 +1,32 @@
 [deps]
+AlgebraOfGraphics = "cbdf2221-f076-402e-a563-3d30da359d67"
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 CondaPkg = "992eb4ea-22a4-4c89-a5bb-47a3300528ab"
 DataDeps = "124859b0-ceae-595e-8997-d05f6a7a8dfe"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
-LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
 MAT = "23992714-dd62-5051-b70f-ba57cb901cac"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 NeuralOperators = "ea5c82af-86e5-48da-8ee1-382d6ad7af4b"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
-Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
+
+[sources]
+NeuralOperators = {path = ".."}
 
 [compat]
-CairoMakie = "0.12.11"
+AlgebraOfGraphics = "0.10.7"
+CairoMakie = "0.13"
 CondaPkg = "0.2.23"
 DataDeps = "0.7.13"
 Documenter = "1.7.0"
 Lux = "1"
-LuxCUDA = "0.3.3"
 MAT = "0.10.7"
 MLUtils = "0.4.4"
-NeuralOperators = "0.5"
-Optimisers = "0.3.3"
+NeuralOperators = "0.6"
+Optimisers = "0.4"
 Printf = "1.10"
 PythonCall = "0.9.23"
-Zygote = "0.6.71"
+Reactant = "0.2.127"

docs/make.jl

@@ -17,8 +17,9 @@ makedocs(;
     format=Documenter.HTML(;
         prettyurls=get(ENV, "CI", "false") == "true",
         canonical="https://docs.sciml.ai/NeuralOperators/stable/",
-        assets=["assets/favicon.ico"]),
-    pages
+        assets=["assets/favicon.ico"],
+    ),
+    pages,
 )
 
 deploydocs(; repo="github.com/SciML/NeuralOperators.jl.git", push_preview=true)

docs/pages.jl

@@ -3,10 +3,13 @@ pages = [
     "Pre-built Models" => [
         "FNO" => "models/fno.md",
         "DeepONet" => "models/deeponet.md",
-        "NOMAD" => "models/nomad.md"
+        "NOMAD" => "models/nomad.md",
     ],
     "Tutorials" => [
-        "Burgers Equation" => "tutorials/burgers.md"
+        "Solving Burgers Equation" => [
+            "DeepONet" => "tutorials/burgers_deeponet.md",
+            "FNO" => "tutorials/burgers_fno.md",
+        ],
     ],
-    "API Reference" => "api.md"
+    "API Reference" => "api.md",
 ]

docs/src/api.md

@@ -21,4 +21,5 @@ SpectralKernel
 
 ```@docs
 NeuralOperators.AbstractTransform
+NeuralOperators.FourierTransform
 ```

docs/src/models/deeponet.md

@@ -11,7 +11,7 @@ u(y) \xrightarrow{\text{branch}} & \; b \\
 & \quad \searrow\\
 &\quad \quad \mathcal{G}_{\theta} u(y) = \sum_k b_k t_k \\
 &  \quad \nearrow \\
-y \; \; \xrightarrow{\text{trunk}} \; \; & t  
+y \; \; \xrightarrow{\text{trunk}} \; \; & t
 \end{align*}
 ```
 
@@ -38,48 +38,65 @@ v(x) = \frac{du}{dx} \quad \forall \; x \in [0, 2\pi], \; \alpha \in [0.5, 1]
 ### Copy-pastable code
 
 ```@example deeponet_tutorial
-using NeuralOperators, Lux, Random, Optimisers, Zygote, CairoMakie
+using NeuralOperators, Lux, Random, Optimisers, Reactant
+
+using CairoMakie, AlgebraOfGraphics
+set_aog_theme!()
+const AoG = AlgebraOfGraphics
 
 rng = Random.default_rng()
+Random.seed!(rng, 1234)
+
+xdev = reactant_device()
 
 eval_points = 1
-data_size = 64
+batch_size = 64
 dim_y = 1
 m = 32
 
 xrange = range(0, 2π; length=m) .|> Float32
-u_data = zeros(Float32, m, data_size)
-α = 0.5f0 .+ 0.5f0 .* rand(Float32, data_size)
+α = 0.5f0 .+ 0.5f0 .* rand(Float32, batch_size)
+
+u_data = zeros(Float32, m, batch_size)
+y_data = rand(rng, Float32, 1, eval_points) .* Float32(2π)
+v_data = zeros(Float32, eval_points, batch_size)
 
-y_data = rand(Float32, 1, eval_points, data_size) .* 2π
-v_data = zeros(Float32, eval_points, data_size)
-for i in 1:data_size
+for i in 1:batch_size
     u_data[:, i] .= sin.(α[i] .* xrange)
-    v_data[:, i] .= -inv(α[i]) .* cos.(α[i] .* y_data[1, :, i])
+    v_data[:, i] .= -inv(α[i]) .* cos.(α[i] .* y_data[1, :])
 end
 
 deeponet = DeepONet(
     Chain(Dense(m => 8, σ), Dense(8 => 8, σ), Dense(8 => 8, σ)),
     Chain(Dense(1 => 4, σ), Dense(4 => 8, σ))
 )
 
-ps, st = Lux.setup(rng, deeponet)
-data = [((u_data, y_data), v_data)]
+ps, st = Lux.setup(rng, deeponet) |> xdev;
+
+u_data = u_data |> xdev;
+y_data = y_data |> xdev;
+v_data = v_data |> xdev;
+data = [((u_data, y_data), v_data)];
 
 function train!(model, ps, st, data; epochs=10)
     losses = []
     tstate = Training.TrainState(model, ps, st, Adam(0.001f0))
     for _ in 1:epochs, (x, y) in data
-
-        _, loss,
-        _, tstate = Training.single_train_step!(AutoZygote(), MSELoss(), (x, y),
-            tstate)
-        push!(losses, loss)
+        (_, loss, _, tstate) = Training.single_train_step!(
+            AutoEnzyme(), MSELoss(), (x, y), tstate; return_gradients=Val(false)
+        )
+        push!(losses, Float32(loss))
     end
     return losses
 end
 
 losses = train!(deeponet, ps, st, data; epochs=1000)
 
-lines(losses)
+draw(
+    AoG.data((; losses, iteration=1:length(losses))) *
+    mapping(:iteration => "Iteration", :losses => "Loss (log10 scale)") *
+    visual(Lines);
+    axis=(; yscale=log10),
+    figure=(; title="Using DeepONet to learn the anti-derivative operator")
+)
 ```