refactor: remove NeuralPDELogging in-favor of extension (#901)

SciML · Oct 16, 2024 · 3085ef5 · 3085ef5
1 parent 4294558
commit 3085ef5
Show file tree

Hide file tree

Showing 14 changed files with 144 additions and 266 deletions.
diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
@@ -1,15 +1,15 @@
 steps:
-  - label: "GPU"
+  - label: "CUDA"
     plugins:
       - JuliaCI/julia#v1:
           version: "1"
       - JuliaCI/julia-test#v1:
-           coverage: false # 1000x slowdown
+           coverage: true
     agents:
       queue: "juliagpu"
       cuda: "*"
     env:
-      GROUP: 'GPU'
+      GROUP: 'CUDA'
       JULIA_PKG_SERVER: "" # it often struggles with our large artifacts
       # SECRET_CODECOV_TOKEN: "..."
     timeout_in_minutes: 240

diff --git a/.github/workflows/CompatHelper.yml b/.github/workflows/CompatHelper.yml
@@ -23,4 +23,4 @@ jobs:
       - name: CompatHelper.main()
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-        run: julia -e 'using CompatHelper; CompatHelper.main(;subdirs=["", "docs", "lib/NeuralPDELogging"])'
+        run: julia -e 'using CompatHelper; CompatHelper.main(;subdirs=["", "docs"])'
diff --git a/.github/workflows/Downgrade.yml b/.github/workflows/Downgrade.yml
@@ -55,7 +55,7 @@ jobs:
           GROUP: ${{ matrix.group }}
       - uses: julia-actions/julia-processcoverage@v1
         with:
-          directories: src,lib/NeuralPDELogging/src
+          directories: src,ext
       - uses: codecov/codecov-action@v4
         with:
           files: lcov.info

diff --git a/.github/workflows/Tests.yml b/.github/workflows/Tests.yml
@@ -41,6 +41,6 @@ jobs:
     uses: "SciML/.github/.github/workflows/tests.yml@v1"
     with:
       group: "${{ matrix.group }}"
-      coverage-directories: "src,lib/NeuralPDELogging/src"
+      coverage-directories: "src,ext"
       julia-version: "${{ matrix.version }}"
     secrets: "inherit"
diff --git a/Project.toml b/Project.toml
@@ -44,6 +44,12 @@ Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 WeightInitializers = "d49dbf32-c5c2-4618-8acc-27bb2598ef2d"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
+[weakdeps]
+TensorBoardLogger = "899adc3e-224a-11e9-021f-63837185c80f"
+
+[extensions]
+NeuralPDETensorBoardLoggerExt = "TensorBoardLogger"
+
 [compat]
 ADTypes = "1.9.0"
 Adapt = "4"
@@ -96,6 +102,7 @@ StochasticDiffEq = "6.69.1"
 SymbolicIndexingInterface = "0.3.31"
 SymbolicUtils = "3.7.2"
 Symbolics = "6.14"
+TensorBoardLogger = "0.1.24"
 Test = "1.10"
 WeightInitializers = "1.0.3"
 Zygote = "0.6.71"
@@ -117,7 +124,8 @@ OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
+TensorBoardLogger = "899adc3e-224a-11e9-021f-63837185c80f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "CUDA", "DiffEqNoiseProcess", "ExplicitImports", "Flux", "LineSearches", "LuxCUDA", "LuxCore", "LuxLib", "MethodOfLines", "OptimizationOptimJL", "OrdinaryDiffEq", "Pkg", "SafeTestsets", "StochasticDiffEq", "Test"]
+test = ["Aqua", "CUDA", "DiffEqNoiseProcess", "ExplicitImports", "Flux", "LineSearches", "LuxCUDA", "LuxCore", "LuxLib", "MethodOfLines", "OptimizationOptimJL", "OrdinaryDiffEq", "Pkg", "SafeTestsets", "StochasticDiffEq", "TensorBoardLogger", "Test"]
diff --git a/ext/NeuralPDETensorBoardLoggerExt.jl b/ext/NeuralPDETensorBoardLoggerExt.jl
@@ -0,0 +1,19 @@
+module NeuralPDETensorBoardLoggerExt
+
+using NeuralPDE: NeuralPDE
+using TensorBoardLogger: TBLogger, log_value
+
+function NeuralPDE.logvector(logger::TBLogger, vector::AbstractVector{<:Real},
+        name::AbstractString, step::Integer)
+    foreach(enumerate(vector)) do (j, v)
+        log_value(logger, "$(name)/$(j)", v; step)
+    end
+end
+
+function NeuralPDE.logscalar(logger::TBLogger, scalar::Real, name::AbstractString,
+        step::Integer)
+    log_value(logger, "$(name)", scalar; step)
+    return nothing
+end
+
+end
diff --git a/lib/NeuralPDELogging/LICENSE b/lib/NeuralPDELogging/LICENSE
diff --git a/lib/NeuralPDELogging/Project.toml b/lib/NeuralPDELogging/Project.toml
diff --git a/lib/NeuralPDELogging/src/NeuralPDELogging.jl b/lib/NeuralPDELogging/src/NeuralPDELogging.jl
diff --git a/lib/NeuralPDELogging/test/adaptive_loss_log_tests.jl b/lib/NeuralPDELogging/test/adaptive_loss_log_tests.jl
diff --git a/lib/NeuralPDELogging/test/runtests.jl b/lib/NeuralPDELogging/test/runtests.jl
diff --git a/src/pinn_types.jl b/src/pinn_types.jl
@@ -10,15 +10,8 @@ end
 
 LogOptions(; log_frequency = 50) = LogOptions(log_frequency)
 
-"""This function is defined here as stubs to be overridden by the subpackage NeuralPDELogging if imported"""
-function logvector(logger, v::AbstractVector{<:Real}, name::AbstractString, step::Integer)
-    nothing
-end
-
-"""This function is defined here as stubs to be overridden by the subpackage NeuralPDELogging if imported"""
-function logscalar(logger, s::Real, name::AbstractString, step::Integer)
-    nothing
-end
+logvector(logger, v::AbstractVector{<:Real}, name::AbstractString, step::Integer) = nothing
+logscalar(logger, s::Real, name::AbstractString, step::Integer) = nothing
 
 """
 An encoding of the test function phi that is used for calculating the PDE