Aerosol activation emulators: GP, ET, and ARG-informed (#349)

add the remaining features from Mikhal's work: GP, ET trainings and training with informed data from ARG

docs/Project.toml

@@ -1,6 +1,6 @@
 [deps]
-ClimaParams = "5c42b081-d73a-476f-9059-fd94b934656c"
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
+ClimaParams = "5c42b081-d73a-476f-9059-fd94b934656c"
 CloudMicrophysics = "6a9e3e04-43cd-43ba-94b9-e8782df3c71b"
 Dierckx = "39dd38d3-220a-591b-8e3c-4c3a8c710a94"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"

docs/src/AerosolActivation.md

@@ -278,3 +278,8 @@ make_ARG_figX(5)
 ![](Abdul-Razzak_and_Ghan_fig_3.svg)
 ![](Abdul-Razzak_and_Ghan_fig_4.svg)
 ![](Abdul-Razzak_and_Ghan_fig_5.svg)
+
+## Aerosol activation prediction with ML emulators
+The CloudMicrophysics package offers an advanced feature for predicting the aerosol activation fraction using machine-learned emulators. Users have access to a function that utilizes pre-trained models for the prediction of activation fraction. Additionally, the package includes tools for training machine learning models, tailored to specific datasets or needs. The training process can benefit from incorporating the ARG activation equation, enhancing the relevance of training data. Emulator training, covering neural networks, Gaussian processes, and EvoTrees, is showcased in `test/aerosol_activation_emulators.jl`.
+
+Using ML emulators for predicting aerosol activation is provided through an extension to the main package. This extension will be loaded with `CloudMicrophysics.jl` if both `MLJ.jl` and  `DataFrames.jl` are loaded by the user as well.

ext/Common.jl

@@ -2,6 +2,10 @@ import CSV
 import DataFrames as DF
 using DataFramesMeta
 
+import CloudMicrophysics.AerosolModel as AM
+import CloudMicrophysics.AerosolActivation as AA
+import Thermodynamics.Parameters as TDP
+
 function get_num_modes(df::DataFrame)
     i = 1
     while true
@@ -68,6 +72,83 @@ function preprocess_aerosol_data(X::DataFrame)
     return X
 end
 
+function get_ARG_act_frac(
+    data_row::NamedTuple,
+    ap::CMP.AerosolActivationParameters,
+    aip::CMP.AirProperties,
+    tps::TDP.ThermodynamicsParameters,
+    FT::DataType,
+)
+
+    num_modes = get_num_modes(data_row)
+    @assert num_modes > 0
+    mode_Ns = []
+    mode_means = []
+    mode_stdevs = []
+    mode_kappas = []
+    w = data_row.velocity
+    T = data_row.initial_temperature
+    p = data_row.initial_pressure
+    for i in 1:num_modes
+        push!(mode_Ns, data_row[Symbol("mode_$(i)_N")])
+        push!(mode_means, data_row[Symbol("mode_$(i)_mean")])
+        push!(mode_stdevs, data_row[Symbol("mode_$(i)_stdev")])
+        push!(mode_kappas, data_row[Symbol("mode_$(i)_kappa")])
+    end
+    ad = AM.AerosolDistribution(
+        Tuple(
+            AM.Mode_κ(
+                mode_means[i],
+                mode_stdevs[i],
+                mode_Ns[i],
+                FT(1),
+                FT(1),
+                FT(0),
+                FT(mode_kappas[i]),
+                1,
+            ) for i in 1:num_modes
+        ),
+    )
+    pv0 = TD.saturation_vapor_pressure(tps, FT(T), TD.Liquid())
+    vapor_mix_ratio = pv0 / TD.Parameters.molmass_ratio(tps) / (p - pv0)
+    q_vap = vapor_mix_ratio / (vapor_mix_ratio + 1)
+    q = TD.PhasePartition(FT(q_vap), FT(0), FT(0))
+
+    return collect(
+        AA.N_activated_per_mode(ap, ad, aip, tps, FT(T), FT(p), FT(w), q),
+    ) ./ mode_Ns
+end
+
+function get_ARG_act_frac(
+    X::DataFrame,
+    ap::CMP.AerosolActivationParameters,
+    aip::CMP.AirProperties,
+    tps::TDP.ThermodynamicsParameters,
+    FT::DataType,
+)
+    return transpose(
+        hcat(get_ARG_act_frac.(NamedTuple.(eachrow(X)), ap, aip, tps, FT)...),
+    )
+end
+
+function preprocess_aerosol_data_with_ARG_act_frac(
+    X::DataFrame,
+    ap::CMP.AerosolActivationParameters,
+    aip::CMP.AirProperties,
+    tps::TDP.ThermodynamicsParameters,
+    FT::DataType,
+)
+    f(y) = get_ARG_act_frac(y, ap, aip, tps, FT)
+    num_modes = get_num_modes(X)
+    X = DF.transform(
+        X,
+        AsTable(All()) =>
+            ByRow(x -> f(x)) =>
+                [Symbol("mode_$(i)_ARG_act_frac") for i in 1:num_modes],
+    )
+    return preprocess_aerosol_data(X)
+end
+
 function target_transform(act_frac)
     return @. atanh(2.0 * 0.99 * (act_frac - 0.5))
 end

ext/EmulatorModelsExt.jl

@@ -9,6 +9,22 @@ import CloudMicrophysics.AerosolActivation as AA
 import CloudMicrophysics.AerosolModel as AM
 import CloudMicrophysics.Parameters as CMP
 
+"""
+    N_activated_per_mode(machine, ap, ad, aip, tps, T, p, w, q)
+
+  - `machine` - ML model
+  - `ap`  - a struct with aerosol activation parameters
+  - `ad`  - aerosol distribution struct
+  - `aip` - a struct with air parameters
+  - `tps` - a struct with thermodynamics parameters
+  - `T`   - air temperature
+  - `p`   - air pressure
+  - `w`   - vertical velocity
+  - `q`   - phase partition
+
+Returns the number of activated aerosol particles
+in each aerosol size distribution mode by using a trained emulator.
+"""
 function AA.N_activated_per_mode(
     machine::MLJ.Machine,
     ap::CMP.AerosolActivationParameters,

test/Project.toml

@@ -8,13 +8,16 @@ ClimaParams = "5c42b081-d73a-476f-9059-fd94b934656c"
 CloudMicrophysics = "6a9e3e04-43cd-43ba-94b9-e8782df3c71b"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 DataFramesMeta = "1313f7d8-7da2-5740-9ea0-a2ca25f37964"
+EvoTrees = "f6006082-12f8-11e9-0c9c-0d5d367ab1e5"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
+GaussianProcesses = "891a1506-143c-57d2-908e-e1f8e92e6de9"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 MLJ = "add582a8-e3ab-11e8-2d5e-e98b27df1bc7"
 MLJFlux = "094fc8d1-fd35-5302-93ea-dabda2abf845"
 MLJModels = "d491faf4-2d78-11e9-2867-c94bc002c0b7"
 RootSolvers = "7181ea78-2dcb-4de3-ab41-2b8ab5a31e74"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
+StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Thermodynamics = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c"