diff --git a/projects/ExpoHybrid/ExpoHybridEstim.jl b/projects/ExpoHybrid/ExpoHybridEstim.jl index 476a3050..e75bfad8 100644 --- a/projects/ExpoHybrid/ExpoHybridEstim.jl +++ b/projects/ExpoHybrid/ExpoHybridEstim.jl @@ -96,7 +96,7 @@ hybrid_model = constructHybridModel( input_batchnorm = true ) -out = train(hybrid_model, df, (:k,); nepochs=300, batchsize=64, opt=AdamW(0.01, (0.9, 0.999), 0.01), loss_types=[:mse, :nse], training_loss=:nse, random_seed=123, yscale = identity, monitor_names=[:Resp0, :k]) +out = train(hybrid_model, df, (:k,); nepochs=300, batchsize=64, opt=AdamW(0.01, (0.9, 0.999), 0.01), loss_types=[:mse, :nse], training_loss=:nse, random_seed=123, yscale = identity, monitor_names=[:Resp0, :k], normalize_predictors=:robust) EasyHybrid.poplot(out) diff --git a/src/models/GenericHybridModel.jl b/src/models/GenericHybridModel.jl index 06b27f50..5beac44d 100644 --- a/src/models/GenericHybridModel.jl +++ b/src/models/GenericHybridModel.jl @@ -291,9 +291,9 @@ end # ─────────────────────────────────────────────────────────────────────────── # Forward pass for SingleNNHybridModel (optimized, no branching) -function (m::SingleNNHybridModel)(ds_k, ps, st) - # 1) get features - predictors = ds_k(m.predictors) +function (m::SingleNNHybridModel)(data, ps, st) + + predictors, forcing = data parameters = m.parameters @@ -338,7 +338,7 @@ function (m::SingleNNHybridModel)(ds_k, ps, st) end # 5) unpack forcing data - forcing_data = unpack_keyedarray(ds_k, m.forcing) + forcing_data = unpack_keyedarray(forcing) # 6) merge all parameters all_params = merge(scaled_nn_params, global_params, fixed_params) @@ -354,7 +354,9 @@ function (m::SingleNNHybridModel)(ds_k, ps, st) end # Forward pass for MultiNNHybridModel (optimized, no branching) -function (m::MultiNNHybridModel)(ds_k, ps, st) +function (m::MultiNNHybridModel)(data, ps, st) + + predictors, forcing = data parameters = m.parameters @@ -374,8 +376,8 @@ function (m::MultiNNHybridModel)(ds_k, ps, st) nn_states = NamedTuple() for (nn_name, nn) in pairs(m.NNs) - predictors = m.predictors[nn_name] - nn_out, st_nn = LuxCore.apply(nn, ds_k(predictors), ps[nn_name], st[nn_name]) + predictor_names = m.predictors[nn_name] + nn_out, st_nn = LuxCore.apply(nn, predictors(predictor_names), ps[nn_name], st[nn_name]) nn_outputs = merge(nn_outputs, NamedTuple{(nn_name,), Tuple{typeof(nn_out)}}((nn_out,))) nn_states = merge(nn_states, NamedTuple{(nn_name,), Tuple{typeof(st_nn)}}((st_nn,))) end @@ -412,8 +414,7 @@ function (m::MultiNNHybridModel)(ds_k, ps, st) all_params = merge(scaled_nn_params, global_params, fixed_params) # 6) unpack forcing data - - forcing_data = unpack_keyedarray(ds_k, m.forcing) + forcing_data = unpack_keyedarray(forcing) all_kwargs = merge(forcing_data, all_params) # 7) Apply mechanistic model diff --git a/src/train.jl b/src/train.jl index d392e90f..64d9431b 100644 --- a/src/train.jl +++ b/src/train.jl @@ -1,5 +1,7 @@ export train, TrainResults +using Statistics: mean, std, median + # beneficial for plotting based on type TrainResults? struct TrainResults train_history @@ -41,7 +43,7 @@ Train a hybrid model using the provided data and save the training process to a """ function train(hybridModel, data, save_ps; nepochs=200, batchsize=10, opt=Adam(0.01), patience=typemax(Int), file_name=nothing, loss_types=[:mse, :r2], training_loss=:mse, agg=sum, train_from = nothing, - random_seed=nothing, shuffleobs = false, yscale=log10, monitor_names=[], return_model=:best) + random_seed=nothing, shuffleobs = false, yscale=log10, monitor_names=[], return_model=:best, normalize_predictors=false) #! check if the EasyHybridMakie extension is loaded. ext = Base.get_extension(@__MODULE__, :EasyHybridMakie) if ext === nothing @@ -57,6 +59,13 @@ function train(hybridModel, data, save_ps; nepochs=200, batchsize=10, opt=Adam(0 # ? split training and validation data (x_train, y_train), (x_val, y_val) = splitobs(data_; at=0.8, shuffle=shuffleobs) + + if normalize_predictors !== false + x_train1, x_val1 = scale(x_train[1], x_val[1], normalize_predictors) + x_train = (x_train1, x_train[2:end]...) + x_val = (x_val1, x_val[2:end]...) + end + train_loader = DataLoader((x_train, y_train), batchsize=batchsize, shuffle=true); if isnothing(train_from) @@ -373,114 +382,115 @@ Utility function to see if the data is already in the expected format or if furt # Arguments: - hm: The Hybrid Model -- data: either a Tuple of KeyedArrays or a single KeyedArray. +- data: either a Tuple of KeyedArrays, a single KeyedArray, or a DataFrame. Returns a tuple of KeyedArrays """ -function prepare_data(hm, data::KeyedArray) - targets = hm.targets - predictors_forcing = Symbol[] - - # Collect all predictors and forcing variables by checking property names - for prop in propertynames(hm) - if occursin("predictors", string(prop)) - val = getproperty(hm, prop) - if isa(val, AbstractVector) - append!(predictors_forcing, val) - elseif isa(val, Union{NamedTuple, Tuple}) - append!(predictors_forcing, unique(vcat(values(val)...))) - end +function _collect_predictors_and_forcing(hm) + predictors = Symbol[] + forcing = Symbol[] + for prop in propertynames(hm) + prop_str = string(prop) + val = getproperty(hm, prop) + if occursin("predictors", prop_str) + if isa(val, AbstractVector) + append!(predictors, val) + elseif isa(val, Union{NamedTuple, Tuple}) + append!(predictors, unique(vcat(values(val)...))) end - end - for prop in propertynames(hm) - if occursin("forcing", string(prop)) - val = getproperty(hm, prop) - if isa(val, AbstractVector) - append!(predictors_forcing, val) - elseif isa(val, Union{Tuple, NamedTuple}) - append!(predictors_forcing, unique(vcat(values(val)...))) - end + elseif occursin("forcing", prop_str) + if isa(val, AbstractVector) + append!(forcing, val) + elseif isa(val, Union{NamedTuple, Tuple}) + append!(forcing, unique(vcat(values(val)...))) end end - predictors_forcing = unique(predictors_forcing) - - if isempty(predictors_forcing) - @warn "Note that you don't have predictors or forcing variables." - end - if isempty(targets) - @warn "Note that you don't have target names." - end - return (data(predictors_forcing), data(targets)) end + return (unique(predictors), unique(forcing)) +end - function prepare_data(hm, data::DataFrame) - targets = hm.targets - predictors_forcing = Symbol[] - - # Collect all predictors and forcing variables by checking property names - for prop in propertynames(hm) - if occursin("predictors", string(prop)) - val = getproperty(hm, prop) - if isa(val, AbstractVector) - append!(predictors_forcing, val) - elseif isa(val, Union{NamedTuple, Tuple}) - append!(predictors_forcing, unique(vcat(values(val)...))) - end - end - end - for prop in propertynames(hm) - if occursin("forcing", string(prop)) - val = getproperty(hm, prop) - if isa(val, AbstractVector) - append!(predictors_forcing, val) - elseif isa(val, Union{Tuple, NamedTuple}) - append!(predictors_forcing, unique(vcat(values(val)...))) - end - end - end - predictors_forcing = unique(predictors_forcing) - - if isempty(predictors_forcing) - @warn "Note that you don't have predictors or forcing variables." - end - if isempty(targets) - @warn "Note that you don't have target names." - end +function _check_predictors_forcing_targets(predictors, forcing, targets) + isempty(predictors) && @warn "No predictors variables." + isempty(forcing) && @warn "No forcing variables - is a hybrid model?" + isempty(targets) && @warn "No target names." +end - all_predictor_cols = unique(vcat(values(predictors_forcing)...)) - col_to_select = unique([all_predictor_cols; targets]) - - # subset to only the cols we care about - sdf = data[!, col_to_select] - - # Separate predictor/forcing vs. target columns - predforce_cols = setdiff(col_to_select, targets) - - # For each row, check if *any* predictor/forcing is missing - mask_missing_predforce = map(row -> any(ismissing, row), eachrow(sdf[:, predforce_cols])) - - # For each row, check if *at least one* target is present (i.e. not all missing) - mask_at_least_one_target = map(row -> any(!ismissing, row), eachrow(sdf[:, targets])) - - # Keep rows where predictors/forcings are *complete* AND there's some target present - keep = .!mask_missing_predforce .& mask_at_least_one_target - sdf = sdf[keep, col_to_select] - - mapcols(col -> replace!(col, missing => NaN), sdf; cols = names(sdf, Union{Missing, Real})) - - # Convert to Float32 and to your keyed array - ds_keyed = to_keyedArray(Float32.(sdf)) - return prepare_data(hm, ds_keyed) - end +function prepare_data(hm, data::KeyedArray) + targets = hm.targets + predictors, forcing = _collect_predictors_and_forcing(hm) + _check_predictors_forcing_targets(predictors, forcing, targets) + return prepare_data(hm, data, predictors, forcing, targets) +end + +function prepare_data(hm, data::DataFrame) + targets = hm.targets + predictors, forcing = _collect_predictors_and_forcing(hm) + _check_predictors_forcing_targets(predictors, forcing, targets) + + all_predictor_cols = unique(vcat(values(predictors)...)) + all_forcing_cols = unique(vcat(values(forcing)...)) + col_to_select = unique([all_predictor_cols; all_forcing_cols; targets]) + + # subset to only the cols we care about + sdf = data[!, col_to_select] + + # Separate predictor/forcing vs. target columns + predforce_cols = setdiff(col_to_select, targets) + + # For each row, check if *any* predictor/forcing is missing + mask_missing_predforce = map(row -> any(ismissing, row), eachrow(sdf[:, predforce_cols])) + + # For each row, check if *at least one* target is present (i.e. not all missing) + mask_at_least_one_target = map(row -> any(!ismissing, row), eachrow(sdf[:, targets])) + + # Keep rows where predictors/forcings are *complete* AND there's some target present + keep = .!mask_missing_predforce .& mask_at_least_one_target + sdf = sdf[keep, col_to_select] + + mapcols(col -> replace!(col, missing => NaN), sdf; cols = names(sdf, Union{Missing, Real})) + + # Convert to Float32 and to your keyed array + ds_keyed = to_keyedArray(Float32.(sdf)) + return prepare_data(hm, ds_keyed, predictors, forcing, targets) +end function prepare_data(hm, data::Tuple) return data end +function prepare_data(hm::Union{SingleNNHybridModel, MultiNNHybridModel}, data::KeyedArray, predictors, forcing, targets) + return ((data(predictors), data(forcing)), data(targets)) +end + +function prepare_data(hm::LuxCore.AbstractLuxContainerLayer, data::KeyedArray, predictors, forcing, targets) + predictors_forcing = unique(vcat(predictors, forcing)) + return (data(predictors_forcing), data(targets)) +end + function get_ps_st(train_from::TrainResults) return train_from.ps, train_from.st end function get_ps_st(train_from::Tuple) return train_from +end + +function scale(x_train::KeyedArray, x_val::KeyedArray, method::Symbol) + if method == :zscore + row_center = mean(x_train, dims=:col) + row_scale = std(x_train, dims=:col) + elseif method == :minmax + row_center = minimum(x_train, dims=:col) + row_scale = maximum(x_train, dims=:col) - row_center + elseif method == :robust + row_center = median(x_train, dims=:col) + row_scale = mapslices(s -> quantile(vec(s), eltype(row_center)(0.75)) - quantile(vec(s), eltype(row_center)(0.25)), + x_train; dims = :col) + end + + row_scale[row_scale .== 0] .= eltype(row_scale)(1e-6) + scaled_train = (x_train .- row_center) ./ row_scale + scaled_val = (x_val .- row_center) ./ row_scale + + return scaled_train, scaled_val end \ No newline at end of file