Merge pull request #23 from CosmologicalEmulators/develop

Develop

Author: marcobonici

Project.toml

@@ -1,7 +1,7 @@
 name = "Capse"
 uuid = "994f66c3-d2c7-4ba6-88fb-4a10f50800ba"
 authors = ["marcobonici <bonici.marco@gmail.com>"]
-version = "0.2.3"
+version = "0.3.0"
 
 [deps]
 AbstractCosmologicalEmulators = "c83c1981-e5c4-4837-9eb8-c9b1572acfc6"
@@ -10,7 +10,7 @@ JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
 NPZ = "15e1cf62-19b3-5cfa-8e77-841668bca605"
 
 [compat]
-AbstractCosmologicalEmulators = "0.3.3"
+AbstractCosmologicalEmulators = "0.4"
 Adapt = "3"
 JSON = "0.21"
 NPZ = "0.4"

docs/Project.toml

@@ -4,3 +4,4 @@ Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
 NPZ = "15e1cf62-19b3-5cfa-8e77-841668bca605"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
+SimpleChains = "de6bee2f-e2f4-4ec7-b6ed-219cc6f6e9e5"

docs/src/assets/postprocessing.jl

@@ -0,0 +1,3 @@
+function postprocessing(input, output, Cℓemu)
+    return output .* exp(input[1]-3.)
+end

docs/src/index.md

@@ -7,6 +7,18 @@ using JSON
 using BenchmarkTools
 using NPZ
 using Capse
+using SimpleChains
+
+mlpd = SimpleChain(
+  static(6),
+  TurboDense(tanh, 64),
+  TurboDense(tanh, 64),
+  TurboDense(tanh, 64),
+  TurboDense(tanh, 64),
+  TurboDense(tanh, 64),
+  TurboDense(identity, 4999)
+)
+
 default(palette = palette(:tab10))
 benchmark = BenchmarkTools.load("./assets/capse_benchmark.json")
 path_json = "./assets/nn_setup.json"
@@ -15,15 +27,16 @@ weights = rand(500000)
 ℓgrid = ones(2000)
 InMinMax_array = zeros(2,2000)
 OutMinMax_array = zeros(2,2000)
-npzwrite("./assets/l.npy", ℓgrid)
-npzwrite("./assets/weights.npy", weights)
-npzwrite("./assets/inminmax.npy", InMinMax_array)
-npzwrite("./assets/outminmax.npy", OutMinMax_array)
-weights_folder = "./assets/"
-Cℓ_emu = Capse.load_emulator(weights_folder)
+nn_setup = JSON.parsefile(path_json)
+emu = Capse.SimpleChainsEmulator(Architecture = mlpd, Weights = weights,
+                                 Description = nn_setup)
+postprocessing(input, output, Cℓemu) = output .* exp(input[1]-3.)
+Cℓ_emu = Capse.CℓEmulator(TrainedEmulator = emu, ℓgrid=ℓgrid, InMinMax = InMinMax_array,
+                                OutMinMax = OutMinMax_array,
+                                Postprocessing = postprocessing)
 ```
 
-`Capse.jl` is a Julia package designed to emulate the computation of the CMB Angular Power Spectrum, with a speedup of several orders of magnitude.
+`Capse.jl` is a Julia package designed to emulate the computation of the CMB Angular Power Spectrum, with a speedup of several orders of magnitude compared to standard codes.
 
 ## Installation
 
@@ -40,7 +53,7 @@ In order to be able to use `Capse.jl`, there two major steps that need to be per
 - Instantiating the emulators, e.g. initializing the Neural Network, its weight and biases and the quantities employed in pre and post-processing
 - Use the instantiated emulators to retrieve the spectra
 
-In the reminder of this section we are showing how this can be done.
+In the reminder of this section we are showing how to do this.
 
 ### Instantiation
 
@@ -52,10 +65,10 @@ Cℓ_emu = Capse.load_emulator(weights_folder);
 
 where `weights_folder` is the path to the folder containing the files required to build up the network. Some of the trained emulators can be found on [Zenodo](https://zenodo.org/record/8187935) and we plan to release more of them there in the future.
 
-It is possible to pass an additional argument to the previous function, which is used to choose between the two NN backed now available:
+It is possible to pass an additional argument to the previous function, which is used to choose between the two NN backend now available:
 
 - [SimpleChains](https://github.com/PumasAI/SimpleChains.jl), which is taylored for small NN running on a CPU
-- [Lux](https://github.com/LuxDL/Lux.jl), which can run on a GPU
+- [Lux](https://github.com/LuxDL/Lux.jl), which can run both on CPUs and GPUs
 
 `SimpleChains.jl` is faster expecially for small NN on the CPU. If you wanna use something running on a GPU, you should use `Lux.jl`, which can be done adding an additional argument to the `load_emulator` function, `Capse.LuxEmulator`
 
@@ -69,17 +82,20 @@ Each trained emulator should be shipped with a description within the JSON file.
 Capse.get_emulator_description(Cℓ_emu)
 ```
 
-After loading a trained emulator, feed it some input parameters `x` in order to get the emulated $C_\ell$'s
+!!! warning
+
+    Cosmological parameters must be fed to `Capse.jl` with **arrays**. It is the user
+    responsability to check the right ordering, by reading the output of the
+    `get_emulator_description` method.
+
+After loading a trained emulator, feed it some input parameters `x` in order to get the
+emulated $C_\ell$'s
 
 ```julia
 x = rand(6) # generate some random input
 Capse.get_Cℓ(x, Cℓ_emu) #compute the Cℓ's
 ```
 
-!!! warning
-
-    In this moment the API is **not** stable: we need to pass the input cosmological parameters in an hardcoded way. We are working to add a more stable and flexible API.
-
 Using `SimpleChains.jl`, we obtain a mean execution time of 45 microseconds
 
 ```@example tutorial
@@ -94,24 +110,26 @@ benchmark[1]["Capse"]["Lux"] # hide
 
 `SimpleChains.jl` is about 2 times faster than `Lux.jl` and they give the same result up to floating point precision.
 
-This benchmarks have been performed locally, with a 12th Gen Intel® Core™ i7-1260P.
+These benchmarks have been performed locally, with a 12th Gen Intel® Core™ i7-1260P.
 
 Considering that a high-precision settings calculation performed with [`CAMB`](https://github.com/cmbant/CAMB) on the same machine requires around 60 seconds, `Capse.jl` is 5-6 order of magnitudes faster.
 
 ### Authors
 
-- Marco Bonici, INAF - Institute of Space Astrophysics and Cosmic Physics (IASF), Milano
+- Marco Bonici, PostDoctoral researcher at Waterloo Center for Astrophysics
 - Federico Bianchini, PostDoctoral researcher at Stanford
 - Jaime Ruiz-Zapatero, PhD Student at Oxford
+- Marius Millea, Researcher at UC Davis and Berkeley Center for Cosmological Physics
 
 ## Contributing
 
-Pull requests are welcome. For major changes, please open an issue first to discuss what you would like to change.
+Pull requests are welcome. For major changes, please open an issue first to discuss what you
+would like to change.
 
 Please make sure to update tests as appropriate.
 
 ### License
 
 `Capse.jl` is licensed under the MIT "Expat" license; see
-[LICENSE](https://github.com/CosmologicalEmulators/Effort.jl/blob/main/LICENSE) for
-the full license text.
+[LICENSE](https://github.com/CosmologicalEmulators/Effort.jl/blob/main/LICENSE) for the full
+license text.

src/Capse.jl

@@ -24,13 +24,16 @@ It contains:
 
 - `InMinMax::AbstractMatrix`, the `Matrix` used for the MinMax normalization of the input features
 
-- O`utMinMax::AbstractMatrix`, the `Matrix` used for the MinMax normalization of the output features
+- `OutMinMax::AbstractMatrix`, the `Matrix` used for the MinMax normalization of the output features
+
+- `Postprocessing::Function`, the `Function` used for the postprocessing of the NN output
 """
 @kwdef mutable struct CℓEmulator <: AbstractCℓEmulators
     TrainedEmulator::AbstractTrainedEmulators
     ℓgrid::AbstractVector
     InMinMax::AbstractMatrix
     OutMinMax::AbstractMatrix
+    Postprocessing::Function
 end
 
 Adapt.@adapt_structure CℓEmulator
@@ -45,7 +48,7 @@ function get_Cℓ(input_params, Cℓemu::AbstractCℓEmulators)
     maximin_input!(input, Cℓemu.InMinMax)
     output = Array(run_emulator(input, Cℓemu.TrainedEmulator))
     inv_maximin_output!(output, Cℓemu.OutMinMax)
-    return output .* exp(input_params[1]-3.)
+    return Cℓemu.Postprocessing(input_params, output, Cℓemu)
 end
 
 """
@@ -73,19 +76,25 @@ The following keyword arguments are used to specify the name of the files used t
 - `inminmax_file`, default `inminmax.npy`
 - `outminmax_file`, default `outminmax.npy`
 - `nn_setup_file`, default `nn_setup.json`
-If the corresponding file in the folder you are trying to load have different names, just change the default values.
+- `postprocessing_file`, default `postprocessing.jl`
+If the corresponding file in the folder you are trying to load have different names,
+ change the default values accordingly.
 """
 function load_emulator(path::String, emu = SimpleChainsEmulator,
     ℓ_file = "l.npy", weights_file = "weights.npy", inminmax_file = "inminmax.npy",
-    outminmax_file = "outminmax.npy", nn_setup_file = "nn_setup.json")
+    outminmax_file = "outminmax.npy", nn_setup_file = "nn_setup.json",
+    postprocessing_file = "postprocessing.jl")
     NN_dict = parsefile(path*nn_setup_file)
     ℓ = npzread(path*ℓ_file)
+    include(path*postprocessing_file)
+    #we assume there is a postprocessing() function in the postprocessing_file
 
     weights = npzread(path*weights_file)
     trained_emu = Capse.init_emulator(NN_dict, weights, emu)
     Cℓ_emu = Capse.CℓEmulator(TrainedEmulator = trained_emu, ℓgrid = ℓ,
                              InMinMax = npzread(path*inminmax_file),
-                             OutMinMax = npzread(path*outminmax_file))
+                             OutMinMax = npzread(path*outminmax_file),
+                             Postprocessing = postprocessing)
     return Cℓ_emu
 end
 

test/emu/postprocessing.jl

@@ -0,0 +1,3 @@
+function postprocessing(input, output, Cℓemu)
+    return output .* exp(input[1]-3.)
+end

test/runtests.jl

@@ -24,8 +24,10 @@ npzwrite("emu/inminmax.npy", inminmax)
 npzwrite("emu/outminmax.npy", outminmax)
 emu = Capse.SimpleChainsEmulator(Architecture = mlpd, Weights = weights)
 
+postprocessing(input, output, Cℓemu) = output .* exp(input[1]-3.)
+
 capse_emu = Capse.CℓEmulator(TrainedEmulator = emu, ℓgrid=ℓ_test, InMinMax = inminmax,
-                                OutMinMax = outminmax)
+                                OutMinMax = outminmax, Postprocessing = postprocessing)
 capse_loaded_emu = Capse.load_emulator("emu/")
 
 @testset "Capse tests" begin
@@ -35,6 +37,6 @@ capse_loaded_emu = Capse.load_emulator("emu/")
     output_vec = Capse.get_Cℓ(cosmo_vec, capse_emu)
     @test isapprox(output_vec[:,1], output)
     @test ℓ_test == Capse.get_ℓgrid(capse_emu)
-    @test_logs (:warn, "We do not know which parameters were included in the emulators training space. Use this trained emulator with caution!") Capse.get_emulator_description(capse_emu)
+    @test_logs (:warn, "No emulator description present!") Capse.get_emulator_description(capse_emu)
     @test Capse.get_Cℓ(cosmo_vec, capse_emu) == Capse.get_Cℓ(cosmo_vec, capse_loaded_emu)
 end