added named targets

Author: lazarusA

projects/RbQ10/Q10.jl

@@ -6,24 +6,27 @@ using Random
 using LuxCore
 using CSV, DataFrames
 using EasyHybrid.MLUtils
-
+using EasyHybrid.AxisKeys
+using Zygote
 # data
-df = CSV.read("/Users/lalonso/Documents/HybridML/data/Rh_AliceHolt_forcing_filled.csv", DataFrame)
+df_o = CSV.read("/Users/lalonso/Documents/HybridML/data/Rh_AliceHolt_forcing_filled.csv", DataFrame)
 
+df = copy(df_o)
 df[!, :Temp] = df[!, :Temp] .- 273.15 # convert to Celsius
-df_forcing = filter(:Respiration_heterotrophic => !isnan, df)
+# df = filter(:Respiration_heterotrophic => !isnan, df)
+rename!(df, :Respiration_heterotrophic => :Rh)  # rename as in hybrid model
 df_forcing = df
-ds_k = to_keyedArray(Float32.(df_forcing))
-yobs =  ds_k(:Respiration_heterotrophic)'[:,:]
+
+ds_p_f = to_keyedArray(Float32.(df_forcing)) # predictors + forcing
+ds_t =  ds_p_f([:Rh]) # do the array so that you conserve the name
 
 NN = Lux.Chain(Dense(2, 15, Lux.relu), Dense(15, 15, Lux.relu), Dense(15, 1));
 #? do different initial Q10s
-RbQ10 = RespirationRbQ10(NN, (:Rgpot, :Moist), (:Temp,), 2.5f0) 
+RbQ10 = RespirationRbQ10(NN, (:Rgpot, :Moist), (:Rh, ), (:Temp,), 2.5f0) 
 
 # ? play with :Temp as predictors in NN, temperature sensitivity!
 # TODO: variance effect due to LSTM vs NN
-
-out = train(RbQ10, (ds_k([:Rgpot, :Moist, :Temp]), yobs), (:Q10, ); nepochs=1000, batchsize=512, opt=Adam(0.01));
+out = train(RbQ10, (ds_p_f, ds_t), (:Q10, ); nepochs=1000, batchsize=512, opt=Adam(0.01));
 
 
 with_theme(theme_light()) do 
@@ -39,10 +42,10 @@ with_theme(theme_light()) do
     fig = Figure(; size = (1200, 600))
     ax_train = Makie.Axis(fig[1, 1], title = "training")
     ax_val = Makie.Axis(fig[2, 1], title = "validation")
-    lines!(ax_train, out.ŷ_train[:], color=:orangered, label = "prediction")
+    lines!(ax_train, out.ŷ_train.Rh[:], color=:orangered, label = "prediction")
     lines!(ax_train, out.y_train[:], color=:dodgerblue, label ="observation")
     # validation
-    lines!(ax_val, out.ŷ_val[:], color=:orangered, label = "prediction")
+    lines!(ax_val, out.ŷ_val.Rh[:], color=:orangered, label = "prediction")
     lines!(ax_val, out.y_val[:], color=:dodgerblue, label ="observation")
     axislegend(; position=:lt)
     Label(fig[0,1], "Observations vs predictions", tellwidth=false)
@@ -62,22 +65,20 @@ with_theme(theme_light()) do
 end
 
 
-ds_k = to_keyedArray(Float32.(df))
-yobs_all =  ds_k(:Respiration_heterotrophic)'[:,:]
+yobs_all =  ds_p_f(:Rh)
 
-ŷ, RbQ10_st = LuxCore.apply(RbQ10, ds_k, out.ps, out.st)
+ŷ, RbQ10_st = LuxCore.apply(RbQ10, ds_p_f, out.ps, out.st)
 
 with_theme(theme_light()) do 
     fig = Figure(; size = (1200, 600))
     ax_train = Makie.Axis(fig[1, 1], title = "full time series")
-    lines!(ax_train, ŷ[:], color=:orangered, label = "prediction")
+    lines!(ax_train, ŷ.Rh[:], color=:orangered, label = "prediction")
     lines!(ax_train, yobs_all[:], color=:dodgerblue, label ="observation")
     axislegend(ax_train; position=:lt)
     Label(fig[0,1], "Observations vs predictions", tellwidth=false)
     fig
 end
 
-
 # ? Rb
 lines(out.αst_train.Rb[:])
-lines!(ds_k(:Moist)[:])
+lines!(ds_p_f(:Moist)[:])

src/models/Respiration_Rb_Q10.jl

@@ -2,17 +2,18 @@
 export RespirationRbQ10
 
 """
-    RespirationRbQ10(NN, predictors, forcing, Q10)
+    RespirationRbQ10(NN, predictors, forcing, targets, Q10)
 
-A linear hybrid model with a neural network `NN`, `predictors` and `forcing` terms.
+A linear hybrid model with a neural network `NN`, `predictors`, `targets` and `forcing` terms.
 """
-struct RespirationRbQ10{D, T1, T2, T3} <: LuxCore.AbstractLuxContainerLayer{(:NN, :predictors, :forcing, :Q10)}
+struct RespirationRbQ10{D, T1, T2, T3, T4} <: LuxCore.AbstractLuxContainerLayer{(:NN, :predictors, :forcing, :targets, :Q10)}
     NN
     predictors
     forcing
+    targets
     Q10
-    function RespirationRbQ10(NN::D, predictors::T1, forcing::T2, Q10::T3) where {D, T1, T2, T3}
-        new{D, T1, T2, T3}(NN, collect(predictors), collect(forcing), [Q10])
+    function RespirationRbQ10(NN::D, predictors::T1, forcing::T2, targets::T3, Q10::T4) where {D, T1, T2, T3, T4}
+        new{D, T1, T2, T3, T4}(NN, collect(predictors), collect(targets), collect(forcing), [Q10])
     end
 end
 
@@ -28,7 +29,7 @@ function LuxCore.initialstates(::AbstractRNG, layer::RespirationRbQ10)
 end
 
 """
-    RespirationRbQ10(NN, predictors, forcing, Q10)(ds_k)
+    RespirationRbQ10(NN, predictors, forcing, targets, Q10)(ds_k)
 
 # Model definition `ŷ = Rb(αᵢ(t)) * Q10^((T(t) - T_ref)/10)`
 
@@ -37,11 +38,11 @@ ŷ (respiration rate) is computed as a function of the neural network output `R
 """
 function (hm::RespirationRbQ10)(ds_k, ps, st::NamedTuple)
     p = ds_k(hm.predictors)
-    x = ds_k(hm.forcing)
+    x = Array(ds_k(hm.forcing)) # don't propagate names after this
 
     Rb, st = LuxCore.apply(hm.NN, p, ps.ps, st.st) #! NN(αᵢ(t)) ≡ Rb(T(t), M(t))
 
-    ŷ = Rb .* ps.Q10 .^(0.1f0 * (x .- 15.0f0)) # ? should 15°C be the reference temperature also an input variable?
+    Rh = Rb .* ps.Q10 .^(0.1f0 * (x .- 15.0f0)) # ? should 15°C be the reference temperature also an input variable?
 
-    return ŷ, (; Rb, st)
+    return (; Rh), (; Rb, st)
 end

src/train.jl

@@ -12,8 +12,8 @@ function train(hybridModel, data, save_ps; nepochs=200, batchsize=10, opt=Adam(0
     opt_state = Optimisers.setup(opt, ps)
 
     # ? initial losses
-    is_no_nan_t = .!isnan.(y_train)[1, :]
-    is_no_nan_v = .!isnan.(y_val)[1, :]
+    is_no_nan_t = .!isnan.(y_train)
+    is_no_nan_v = .!isnan.(y_val)
     l_init_train = lossfn(hybridModel, x_train, (y_train, is_no_nan_t), ps, st)
     l_init_val = lossfn(hybridModel, x_val, (y_val, is_no_nan_v), ps, st)
 
@@ -24,8 +24,8 @@ function train(hybridModel, data, save_ps; nepochs=200, batchsize=10, opt=Adam(0
     for epoch in 1:nepochs
         for (x, y) in train_loader
             # ? check NaN indices before going forward, and pass filtered `x, y`.
-            is_no_nan = .!isnan.(y)[1, :]
-            if length(is_no_nan)>0
+            is_no_nan = .!isnan.(y)
+            if length(is_no_nan)>0 # ! be careful here, multivariate needs fine tuning
                 grads = Zygote.gradient((ps) -> lossfn(hybridModel, x, (y, is_no_nan), ps, st), ps)[1]
                 Optimisers.update!(opt_state, ps, grads)
             end

src/utils/losses.jl

@@ -12,11 +12,16 @@ end
 
 
 """
-    lossfn(HM::RespirationRbQ10, ds, y, ps, st)
+    lossfn(HM::RespirationRbQ10, ds_p, (ds_t, ds_t_nan), ps, st)
 """
-function lossfn(HM::RespirationRbQ10, ds, (y, no_nan), ps, st)
-    ŷ, αst = HM(ds, ps, st)
-    _, st = αst
-    loss = mean((y[no_nan] .- ŷ[no_nan]).^2)
+function lossfn(HM::RespirationRbQ10, ds_p, (ds_t, ds_t_nan), ps, st)
+    ŷ, _ = HM(ds_p, ps, st)
+    y = ds_t(HM.targets)
+    y_nan = ds_t_nan(HM.targets)
+
+    loss = 0.0
+    for k in axiskeys(y, 1)
+        loss += mean(abs2, (ŷ[k][y_nan(k)] .- y(k)[y_nan(k)]))
+    end
     return loss
 end