diff --git a/Project.toml b/Project.toml index 4432d02..de46eb1 100644 --- a/Project.toml +++ b/Project.toml @@ -1,13 +1,16 @@ name = "Capse" uuid = "994f66c3-d2c7-4ba6-88fb-4a10f50800ba" authors = ["marcobonici "] - -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" diff --git a/README.md b/README.md index 6f4c36f..c920899 100644 --- a/README.md +++ b/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. diff --git a/docs/Project.toml b/docs/Project.toml index 635eb32..f684b6e 100644 --- a/docs/Project.toml +++ b/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" diff --git a/docs/src/api.md b/docs/src/api.md index 9201419..7a63985 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -4,4 +4,6 @@ Capse.CℓEmulator Capse.get_Cℓ Capse.get_ℓgrid +Capse.get_emulator_description +Capse.load_emulator ``` diff --git a/docs/src/index.md b/docs/src/index.md index 587047f..2d57853 100644 --- a/docs/src/index.md +++ b/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 diff --git a/src/Capse.jl b/src/Capse.jl index 88546c1..9bc40f1 100644 --- a/src/Capse.jl +++ b/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 diff --git a/test/Project.toml b/test/Project.toml index eebfd45..abaf678 100644 --- a/test/Project.toml +++ b/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" diff --git a/test/emu/nn_setup.json b/test/emu/nn_setup.json new file mode 100644 index 0000000..712cc8d --- /dev/null +++ b/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" : "bonici.marco@gmail.com", + "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, τ" + } +} diff --git a/test/runtests.jl b/test/runtests.jl index 5a7491c..b7ad6d8 100644 --- a/test/runtests.jl +++ b/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