Merge pull request #80 from CosmologicalEmulators/develop

Update ci.yml

Author: marcobonici

.github/workflows/ci.yml

@@ -12,7 +12,7 @@ on:
     tags: "*"
 jobs:
   test:
-    name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ github.event_name }}
+    name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ github.event_name }}
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
@@ -42,11 +42,15 @@ jobs:
             ${{ runner.os }}-
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
+      # Only process and upload coverage once, on a single representative job
       - uses: julia-actions/julia-processcoverage@v1
+        if: matrix.version == '1.10' && matrix.os == 'ubuntu-latest'
         with:
           directories: src,ext
       - uses: codecov/codecov-action@v5
+        if: matrix.version == '1.10' && matrix.os == 'ubuntu-latest'
         with:
           files: lcov.info
           token: ${{ secrets.CODECOV_TOKEN }}
           fail_ci_if_error: false
+          verbose: true

Project.toml

@@ -1,35 +1,42 @@
 name = "AbstractCosmologicalEmulators"
 uuid = "c83c1981-e5c4-4837-9eb8-c9b1572acfc6"
-version = "0.9.3"
+version = "0.9.4"
 authors = ["Marco Bonici <bonici.marco@gmail.com>"]
 
 [deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 NPZ = "15e1cf62-19b3-5cfa-8e77-841668bca605"
 SimpleChains = "de6bee2f-e2f4-4ec7-b6ed-219cc6f6e9e5"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [weakdeps]
 DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0"
 FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838"
 Integrals = "de52edbc-65ea-441a-8357-d3a637375a31"
-Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
 SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
 
 [extensions]
 BackgroundCosmologyExt = ["DataInterpolations", "FastGaussQuadrature", "Integrals", "LinearAlgebra", "OrdinaryDiffEqTsit5", "SciMLSensitivity"]
-MooncakeExt = ["Mooncake"]
+MooncakeExt = ["ForwardDiff", "Mooncake"]
 
 [compat]
+ADTypes = "1.21.0"
 Adapt = "3, 4"
+BenchmarkTools = "1.6.3"
 ChainRulesCore = "1.26"
 DataInterpolations = "6, 8"
+DifferentiationInterface = "0.7.16"
 FFTW = "1"
 FastGaussQuadrature = "1"
 ForwardDiff = "0.10, 1"
@@ -42,6 +49,7 @@ NPZ = "0.4"
 OrdinaryDiffEqTsit5 = "1"
 SciMLSensitivity = "7.90"
 SimpleChains = "0.4"
+Zygote = "0.7.10"
 julia = "1.10"
 
 [extras]

ext/MooncakeExt/MooncakeExt.jl

@@ -48,4 +48,39 @@ Mooncake.increment_rdata!!(x::FFTW.FFTWPlan, ::NoRData) = x
 
 @from_chainrules MinimalCtx Tuple{typeof(AbstractCosmologicalEmulators.chebyshev_decomposition), Any, Any}
 
+
+using ForwardDiff
+using Lux
+using ChainRulesCore
+
+# Define ChainRulesCore.rrule for run_emulator (specifically for LuxEmulator)
+function ChainRulesCore.rrule(::typeof(AbstractCosmologicalEmulators.run_emulator), input, emulator::AbstractCosmologicalEmulators.LuxEmulator)
+    y = AbstractCosmologicalEmulators.run_emulator(input, emulator)
+    
+    function run_emulator_pullback(Δy)
+        if Δy isa ChainRulesCore.AbstractZero
+            return NoTangent(), ZeroTangent(), NoTangent()
+        end
+        
+        # Ensure Δy is a dense vector
+        Δy_vec = collect(vec(ChainRulesCore.unthunk(Δy)))
+        
+        # ForwardDiff VJP
+        vjp_input = convert(typeof(input), ForwardDiff.gradient(
+            x -> begin
+                y_dual, _ = Lux.apply(emulator.Model, x, emulator.Parameters, emulator.States)
+                sum(y_dual .* Δy_vec)
+            end,
+            input
+        ))
+        
+        return NoTangent(), vjp_input, NoTangent()
+    end
+    
+    return y, run_emulator_pullback
+end
+
+# Register it for Mooncake
+Mooncake.@from_chainrules Mooncake.MinimalCtx Tuple{typeof(AbstractCosmologicalEmulators.run_emulator), Any, AbstractCosmologicalEmulators.LuxEmulator}
+
 end # module MooncakeExt