Merge pull request #20 from CosmologicalEmulators/develop

Develop

Project.toml

@@ -1,13 +1,16 @@
 name = "Capse"
 uuid = "994f66c3-d2c7-4ba6-88fb-4a10f50800ba"
 authors = ["marcobonici <[email protected]>"]
-
-version = "0.2.2"
+version = "0.2.3"
 
 [deps]
 AbstractCosmologicalEmulators = "c83c1981-e5c4-4837-9eb8-c9b1572acfc6"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
+NPZ = "15e1cf62-19b3-5cfa-8e77-841668bca605"
 
 [compat]
 AbstractCosmologicalEmulators = "0.3.3"
 Adapt = "3"
+JSON = "0.21"
+NPZ = "0.4"

README.md

@@ -13,6 +13,34 @@ Repo containing the CMB Angular Power Spectrum Emulator, `Capse.jl`.
 
 `Capse.jl` is entirely written in `Julia`, but can be transparently called from `Python` using the [`pycapse`](https://github.com/CosmologicalEmulators/pycapse) wrapper.
 
+## Installation and usage
+
+Details about installing and using `Capse.jl` can be found in the official [documentation](https://cosmologicalemulators.github.io/Capse.jl/stable/), but can be summerized as follows.
+
+In order to install `Capse.jl`, run from the `Julia` REPL
+
+```julia
+using Pkg, Pkg.add(url="https://github.com/CosmologicalEmulators/Capse.jl")
+```
+
+After installing it, you need to instantiate a trained emulator, which can be done with
+
+```julia
+Cℓ_emu = Capse.load_emulator(weights_folder)
+```
+
+where `weights_folder` is the path to the folder with the trained emulator (some of them can be found on [Zenodo](https://zenodo.org/record/8187935)). After this operation, to obtain the predictions for a set of cosmological parameters, put them in an array `x` and run
+
+```julia
+Capse.get_Cℓ(x, Cℓ_emu)
+```
+
+If you want to retrieve details about the emulators (e.g. which parameters were used to train it etc...) just run
+
+```julia
+Capse.get_emulator_description(Cℓ_emu)
+```
+
 ## Citing
 
 Free usage of the software in this repository is provided, given that you cite our release paper.

docs/Project.toml

@@ -2,4 +2,5 @@
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
+NPZ = "15e1cf62-19b3-5cfa-8e77-841668bca605"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"

docs/src/api.md

@@ -4,4 +4,6 @@
 Capse.CℓEmulator
 Capse.get_Cℓ
 Capse.get_ℓgrid
+Capse.get_emulator_description
+Capse.load_emulator
 ```

docs/src/index.md

@@ -5,14 +5,22 @@ using Plots; gr()
 Plots.reset_defaults()
 using JSON
 using BenchmarkTools
+using NPZ
 using Capse
 default(palette = palette(:tab10))
 benchmark = BenchmarkTools.load("./assets/capse_benchmark.json")
 path_json = "./assets/nn_setup.json"
-weights = rand(20000)
+path_data = "./assets/"
+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)
 ```
 
 `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.
@@ -27,46 +35,32 @@ using Pkg, Pkg.add(url="https://github.com/CosmologicalEmulators/Capse.jl")
 
 ## Usage
 
-If you wanna use `Capse.jl` you need to load a trained emulator.
-We recommend to instantiate a trained `Capse-jl` emulator in the following way.
+In order to be able to use `Capse.jl`, there two major steps that need to be performed:
 
-First of all, instantiate a neural network with its weights and biases. This is done through a method imported through the upstream library [`AbstractCosmologicalEmulators.jl`](https://github.com/CosmologicalEmulators/AbstractCosmologicalEmulators.jl). In order to use it, you need a dictionary containg the information to instantiate the right neural network architecture and an array containing the weights and the biases (both of them can be found on [Zenodo](https://zenodo.org/record/8187935) and we plan to release more of them).
-At the given link above, you can find a `JSON` file with the aforementioned NN architecture. This can be read using the `JSON` library in the following way
+- 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
 
-```@example tutorial
-NN_dict = JSON.parsefile(path_json);
-```
-
-This file contains all the informations required to correctly instantiate the neural network.
-
-After this, you just have to pass the `NN_dict` and the `weights` array to the `init_emulator` method and choose a NN backend. In this moment we support two different Neural Networks libraries:
+In the reminder of this section we are showing how this can be done.
 
-- [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
+### Instantiation
 
-In order to instantiate the emulator, just run
+The most direct way to instantiate an official trained emulators is given by the following one-liner
 
-```@example tutorial
-trained_emu = Capse.init_emulator(NN_dict, weights, Capse.SimpleChainsEmulator);
+```julia
+Cℓ_emu = Capse.load_emulator(weights_folder);
 ```
 
-`SimpleChains.jl` is faster expecially for small NN on the CPU. If you prefer to use `Lux.jl`, pass as last argument `Capse.LuxEmulator`.
+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.
 
-Each trained emulator should be shipped with a description within the JSON file. In order to print the description, just runs:
-
-```@example tutorial
-Capse.get_emulator_description(NN_dict["emulator_description"])
-```
-After instantiating the NN, we need:
+It is possible to pass an additional argument to the previous function, which is used to choose between the two NN backed now available:
 
-- the ``\ell``-grid used to train the emulator, `ℓgrid`
-- the arrays used to perform the minmax normalization of both input and output features, `InMinMax_array` and `OutMinMax_array`
+- [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
 
-Now you can instantiate the emulator, using
+`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`
 
-```@example tutorial
-Cℓ_emu = Capse.CℓEmulator(TrainedEmulator = trained_emu, ℓgrid = ℓgrid,
-                          InMinMax = InMinMax_array, OutMinMax = OutMinMax_array);
+```julia
+Cℓ_emu = Capse.load_emulator(weights_folder, Capse.LuxEmulator);
 ```
 
 Each trained emulator should be shipped with a description within the JSON file. In order to print the description, just runs:
@@ -98,11 +92,11 @@ Using `Lux.jl`, with the same architecture and weights, we obtain
 benchmark[1]["Capse"]["Lux"] # hide
 ```
 
-SimpleChains is about 2 times faster than Lux and they give the same result up to floating point precision.
+`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.
 
-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 around ``1,000,000`` times faster.
+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
 

src/Capse.jl

@@ -4,8 +4,10 @@ using Base: @kwdef
 using Adapt
 using AbstractCosmologicalEmulators
 import AbstractCosmologicalEmulators.get_emulator_description
+import JSON.parsefile
+import NPZ.npzread
 
-export get_Cℓ, get_emulator_description
+export get_Cℓ
 
 abstract type AbstractCℓEmulators end
 
@@ -16,13 +18,13 @@ abstract type AbstractCℓEmulators end
 This is the fundamental struct used to obtain the ``C_\\ell``'s from an emulator.
 It contains:
 
-- TrainedEmulator::AbstractTrainedEmulators, the trained emulator
+- `TrainedEmulator::AbstractTrainedEmulators`, the trained emulator
 
-- ℓgrid::AbstractVector, the ``\\ell``-grid the emulator has been trained on.
+- `ℓgrid::AbstractVector`, the ``\\ell``-grid the emulator has been trained on.
 
-- InMinMax::AbstractMatrix, the `Matrix` used for the MinMax normalization of the input features
+- `InMinMax::AbstractMatrix`, the `Matrix` used for the MinMax normalization of the input features
 
-- OutMinMax::AbstractMatrix, the `Matrix` used for the MinMax normalization of the output features
+- O`utMinMax::AbstractMatrix`, the `Matrix` used for the MinMax normalization of the output features
 """
 @kwdef mutable struct CℓEmulator <: AbstractCℓEmulators
     TrainedEmulator::AbstractTrainedEmulators
@@ -34,14 +36,15 @@ end
 Adapt.@adapt_structure CℓEmulator
 
 """
-    get_Cℓ(CℓEmulator::AbstractCℓEmulators)
-Computes and returns the ``C_\\ell``on the ``\\ell``-grid the emulator has been trained on.
+    get_Cℓ(input_params, Cℓemu::AbstractCℓEmulators)
+Computes and returns the ``C_\\ell``'s on the ``\\ell``-grid the emulator has been trained on given input array `input_params`.
+
 """
-function get_Cℓ(input_params, CℓEmulator::AbstractCℓEmulators)
+function get_Cℓ(input_params, Cℓemu::AbstractCℓEmulators)
     input = deepcopy(input_params)
-    maximin_input!(input, CℓEmulator.InMinMax)
-    output = Array(run_emulator(input, CℓEmulator.TrainedEmulator))
-    inv_maximin_output!(output, CℓEmulator.OutMinMax)
+    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.)
 end
 
@@ -53,8 +56,37 @@ function get_ℓgrid(CℓEmulator::AbstractCℓEmulators)
     return CℓEmulator.ℓgrid
 end
 
-function get_emulator_description(Clemu::AbstractCℓEmulators)
-    get_emulator_description(Clemu.TrainedEmulator)
+"""
+    get_emulator_description(Cℓemu::AbstractCℓEmulators)
+Print on screen the emulator description.
+"""
+function get_emulator_description(Cℓemu::AbstractCℓEmulators)
+    get_emulator_description(Cℓemu.TrainedEmulator)
+end
+
+"""
+    load_emulator(path::String, emu_backend::AbstractTrainedEmulators)
+Load the emulator with the files in the folder `path`, using the backend defined by `emu_backend`.
+The following keyword arguments are used to specify the name of the files used to load the emulator:
+- `ℓ_file`, default `l.npy`
+- `weights_file`, default `weights.npy`
+- `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.
+"""
+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")
+    NN_dict = parsefile(path*nn_setup_file)
+    ℓ = npzread(path*ℓ_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))
+    return Cℓ_emu
 end
 
 end # module

test/Project.toml

@@ -1,5 +1,6 @@
 [deps]
 AbstractCosmologicalEmulators = "c83c1981-e5c4-4837-9eb8-c9b1572acfc6"
+NPZ = "15e1cf62-19b3-5cfa-8e77-841668bca605"
 SimpleChains = "de6bee2f-e2f4-4ec7-b6ed-219cc6f6e9e5"
 Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

test/emu/nn_setup.json

@@ -0,0 +1,33 @@
+{
+    "n_input_features": 6,
+    "layers": {
+        "layer_1": {
+            "n_neurons": 64,
+            "activation_function": "tanh"
+        },
+        "layer_2": {
+            "n_neurons": 64,
+            "activation_function": "tanh"
+        },
+        "layer_3": {
+            "n_neurons": 64,
+            "activation_function": "tanh"
+        },
+        "layer_4": {
+            "n_neurons": 64,
+            "activation_function": "tanh"
+        },
+        "layer_5": {
+            "n_neurons": 64,
+            "activation_function": "tanh"
+        }
+    },
+    "n_output_features": 40,
+    "n_hidden_layers": 5,
+    "emulator_description": {
+        "author" : "Marco Bonici",
+        "author_email" : "[email protected]",
+        "miscellanea" : "The emulator has been trained on the high-precision-settings prediction as computed by the CAMB Boltzmann solver.",
+        "parameters" : "ln10As, ns, H0, ωb, ωc, τ"
+    }
+}

test/runtests.jl

@@ -1,4 +1,5 @@
 using Test
+using NPZ
 using SimpleChains
 using Static
 using Capse
@@ -13,11 +14,19 @@ mlpd = SimpleChain(
   TurboDense(identity, 40)
 )
 
+ℓ_test = Array(LinRange(0,200, 40))
 weights = SimpleChains.init_params(mlpd)
+inminmax = rand(6,2)
+outminmax = rand(40,2)
+npzwrite("emu/l.npy", ℓ_test)
+npzwrite("emu/weights.npy", weights)
+npzwrite("emu/inminmax.npy", inminmax)
+npzwrite("emu/outminmax.npy", outminmax)
 emu = Capse.SimpleChainsEmulator(Architecture = mlpd, Weights = weights)
-ℓ_test = Array(LinRange(0,200, 40))
-capse_emu = Capse.CℓEmulator(TrainedEmulator = emu, ℓgrid=ℓ_test, InMinMax = rand(6,2),
-                                OutMinMax = rand(40,2))
+
+capse_emu = Capse.CℓEmulator(TrainedEmulator = emu, ℓgrid=ℓ_test, InMinMax = inminmax,
+                                OutMinMax = outminmax)
+capse_loaded_emu = Capse.load_emulator("emu/")
 
 @testset "Capse tests" begin
     cosmo = ones(6)
@@ -27,4 +36,5 @@ capse_emu = Capse.CℓEmulator(TrainedEmulator = emu, ℓgrid=ℓ_test, InMinMax
     @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 Capse.get_Cℓ(cosmo_vec, capse_emu) == Capse.get_Cℓ(cosmo_vec, capse_loaded_emu)
 end