Gaussian processes surrogates in FluxMatcher

src/FUSE.jl

@@ -121,6 +121,7 @@ include(joinpath("actors", "divertors", "divertors_actor.jl"))
 include(joinpath("actors", "transport", "neoclassical_actor.jl"))
 include(joinpath("actors", "transport", "tglf_actor.jl"))
 include(joinpath("actors", "transport", "qlgyro_actor.jl"))
+#include(joinpath("actors", "transport", "surrogates_actor.jl"))
 include(joinpath("actors", "transport", "flux_calculator_actor.jl"))
 include(joinpath("actors", "transport", "flux_matcher_actor.jl"))
 include(joinpath("actors", "transport", "eped_profiles_actor.jl"))

src/actors/transport/flux_calculator_actor.jl

@@ -6,7 +6,7 @@ Base.@kwdef mutable struct FUSEparameters__ActorFluxCalculator{T<:Real} <: Param
     _name::Symbol = :not_set
     _time::Float64 = NaN
     rho_transport::Entry{AbstractVector{T}} = Entry{AbstractVector{T}}("-", "rho core transport grid"; default=0.25:0.1:0.85)
-    turbulence_model::Switch{Symbol} = Switch{Symbol}([:TGLF, :QLGYRO, :none], "-", "Turbulence model to use"; default=:TGLF)
+    turbulence_model::Switch{Symbol} = Switch{Symbol}([:TGLF, :QLGYRO,:surrogate, :none], "-", "Turbulence model to use"; default=:TGLF)
     neoclassical_model::Switch{Symbol} = Switch{Symbol}([:neoclassical, :none], "-", "Neocalssical model to use"; default=:neoclassical)
 end
 
@@ -15,6 +15,7 @@ mutable struct ActorFluxCalculator{D,P} <: CompoundAbstractActor{D,P}
     par::FUSEparameters__ActorFluxCalculator{P}
     actor_turb::Union{ActorTGLF{D,P},ActorQLGYRO{D,P},ActorNoOperation{D,P}}
     actor_neoc::Union{ActorNeoclassical{D,P},ActorNoOperation{D,P}}
+    actor_surrogate::Union{ActorSurrogate{D,P},ActorNoOperation{D,P}}
 end
 
 """
@@ -41,6 +42,9 @@ function ActorFluxCalculator(dd::IMAS.dd, par::FUSEparameters__ActorFluxCalculat
     elseif par.turbulence_model == :QLGYRO
         act.ActorQLGYRO.rho_transport = par.rho_transport
         actor_turb = ActorQLGYRO(dd, act.ActorQLGYRO)
+    elseif par.turbulence_model == :surrogate
+        act.ActorSurrogate.generate_surrogates = par.false
+        actor_turb = ActorSurrogate(dd, act.ActorSurrogate)      
     end
 
     if par.neoclassical_model == :none

src/actors/transport/flux_matcher_actor.jl

@@ -23,7 +23,7 @@ Base.@kwdef mutable struct FUSEparameters__ActorFluxMatcher{T<:Real} <: Paramete
     find_widths::Entry{Bool} = Entry{Bool}("-", "Runs turbulent transport actor TJLF finding widths after first iteration"; default=true)
     max_iterations::Entry{Int} = Entry{Int}("-", "Maximum optimizer iterations"; default=500)
     optimizer_algorithm::Switch{Symbol} =
-        Switch{Symbol}([:anderson, :newton, :trust_region, :simple, :none], "-", "Optimizing algorithm used for the flux matching"; default=:anderson)
+        Switch{Symbol}([:anderson, :newton, :trust_region, :simple,:surrogate, :none], "-", "Optimizing algorithm used for the flux matching"; default=:anderson)
     step_size::Entry{T} = Entry{T}(
         "-",
         "Step size for each algorithm iteration (note this has a different meaning for each algorithm)";
@@ -43,6 +43,7 @@ mutable struct ActorFluxMatcher{D,P} <: CompoundAbstractActor{D,P}
     par::FUSEparameters__ActorFluxMatcher{P}
     actor_ct::ActorFluxCalculator{D,P}
     actor_ped::ActorPedestal{D,P}
+    actor_sur::Union{ActorSurrogate{D,P},ActorNoOperation{D,P}}
     norms::Vector{Float64}
     error::Float64
 end
@@ -119,6 +120,8 @@ function _step(actor::ActorFluxMatcher)
             res = (zero=z_init_scaled,)
         elseif par.optimizer_algorithm == :simple
             res = flux_match_simple(actor, z_init_scaled, initial_cp1d, initial_summary_ped, z_scaled_history, err_history, ftol, xtol, prog)
+        if par.optimizer_algorithm == :surrogate
+            res = flux_match_surrogate(actor, z_init_scaled, initial_cp1d, initial_summary_ped, z_scaled_history, err_history, ftol, xtol, prog)
         else
             if par.optimizer_algorithm == :newton
                 opts = Dict(:method => :newton, :factor => par.step_size)
@@ -541,6 +544,7 @@ function flux_match_simple(
     xerror = Inf
     step_size = par.step_size
     max_iterations = par.max_iterations
+
     while (ferror > ftol) || (xerror .> xtol)
         i += 1
         if (i > max_iterations)
@@ -552,11 +556,89 @@ function flux_match_simple(
         targets, fluxes, errors = flux_match_errors(actor, scale_z_profiles(zprofiles), initial_cp1d, initial_summary_ped; z_scaled_history, err_history, prog)
         xerror = maximum(abs.(zprofiles .- zprofiles_old)) / step_size
         zprofiles_old = zprofiles
+
+
     end
 
     return (zero=z_scaled_history[argmin(map(norm, err_history))],)
 end
 
+
+
+"""
+    flux_match_simple(
+        actor::ActorFluxMatcher,
+        z_init::Vector{<:Real},
+        initial_cp1d::IMAS.core_profiles__profiles_1d,
+        initial_summary_ped::IMAS.summary__local__pedestal,
+        z_scaled_history::Vector,
+        err_history::Vector{Float64},
+        ftol::Float64,
+        xtol::Float64,
+        prog::Any)
+
+Updates zprofiles based on TGYRO simple algorithm
+"""
+function flux_match_surrogate(
+    actor::ActorFluxMatcher,
+    z_init_scaled::Vector{<:Real},
+    initial_cp1d::IMAS.core_profiles__profiles_1d,
+    initial_summary_ped::IMAS.summary__local__pedestal,
+    z_scaled_history::Vector,
+    err_history::Vector{Float64},
+    ftol::Float64,
+    xtol::Float64,
+    prog::Any)
+
+    par = actor.par
+    dd = actor.dd
+
+    i = 0
+    zprofiles_old = unscale_z_profiles(z_init_scaled)
+    targets, fluxes, errors = flux_match_errors(actor, z_init_scaled, initial_cp1d, initial_summary_ped; z_scaled_history, err_history, prog)
+    ferror = norm(errors)
+    xerror = Inf
+    step_size = par.step_size
+    max_iterations = par.max_iterations
+
+    all_model_inputs = []
+    all_model_outputs = []
+    while (ferror > ftol) || (xerror .> xtol)
+        i += 1
+        if (i > max_iterations)
+            @info "Unable to flux-match within $(max_iterations) iterations (aerr) = $(ferror) (ftol=$ftol) (xerr) = $(xerror) (xtol = $xtol)"
+            break
+        end
+
+        zprofiles = zprofiles_old .* (1.0 .+ step_size * 0.1 .* (targets .- fluxes) ./ sqrt.(1.0 .+ fluxes .^ 2 + targets .^ 2))
+        targets, fluxes, errors = flux_match_errors(actor, scale_z_profiles(zprofiles), initial_cp1d, initial_summary_ped; z_scaled_history, err_history, prog)
+        xerror = maximum(abs.(zprofiles .- zprofiles_old)) / step_size
+        zprofiles_old = zprofiles
+
+        push!(all_model_inputs, actor.actor_ct.actor_turb.input_tglfs)
+        push!(all_model_outputs, fluxes)
+    end
+
+    max_iterations_surrogates = 1
+    while (ferror > ftol) || (xerror .> xtol)
+        i += 1
+        if (i > max_iterations_surrogates)
+            @info "Unable to flux-match within $(max_iterations) iterations (aerr) = $(ferror) (ftol=$ftol) (xerr) = $(xerror) (xtol = $xtol)"
+            break
+        end
+
+        actor.actor_sur.all_model_inputs = all_model_inputs
+        actor.actor_sur.all_model_outputs = all_model_outputs
+        actor.actor_sur.par.generate_surrogates = true
+        ActorSurrogate(dd, actor.actor_sur)
+
+    end
+
+
+    return (zero=z_scaled_history[argmin(err_history)],)
+end
+
+
 function progress_ActorFluxMatcher(dd::IMAS.dd, error::Float64)
     cp1d = dd.core_profiles.profiles_1d[]
     return (

src/actors/transport/surrogates_actor.jl

@@ -0,0 +1,151 @@
+using Surrogates
+using AbstractGPs #required to access different types of kernels
+using SurrogatesAbstractGPs
+
+#= =========== =#
+#  ActorSurrogate  #
+#= =========== =#
+Base.@kwdef mutable struct FUSEparameters__ActorSurrogate{T<:Real} <: ParametersActor{T}
+    _parent::WeakRef = WeakRef(nothing)
+    _name::Symbol = :not_set
+    _time::Float64 = NaN
+    generate_surrogates::Entry{Bool} = Entry{Bool}("-", "Generate surrogates"; default=false)
+end
+
+mutable struct ActorSurrogate{D,P} <: SingleAbstractActor{D,P}
+    dd::IMAS.dd{D}
+    par::FUSEparameters__ActorSurrogate{P}
+    all_model_inputs::Union{Vector{<:InputTGLF},Any}
+    all_model_outputs::Union{Vector{<:IMAS.flux_solution},Any,Any}
+    flux_surrogates::Union{Vector{<:Any},Any,Any}
+    flux_solutions::Union{Vector{<:IMAS.flux_solution},Any}
+end
+
+"""
+    ActorSurrogate(dd::IMAS.dd, act::ParametersAllActors; kw...)
+
+Evaluates the QLGYRO predicted turbulence
+"""
+function ActorSurrogate(dd::IMAS.dd, act::ParametersAllActors; kw...)
+    actor = ActorSurrogate(dd, act.ActorSurrogate; kw...)
+    step(actor)
+    finalize(actor)
+    return actor
+end
+
+
+struct flux_surrogate{T<:Any}
+    ENERGY_FLUX_e::T
+    ENERGY_FLUX_i::T
+    PARTICLE_FLUX_e::T
+    PARTICLE_FLUX_i::Vector{T}
+    STRESS_TOR_i::T
+end
+
+
+function generate_surrogates(all_input_tglfs, all_flux_solutions)
+
+    actor.all_input_tglfs
+    actor.flux_solutions
+    nsims = length(flux_solutions)
+    nrads = length(flux_solutions[1])
+    flux_surrogates = Vector{flux_surrogate}()
+    for irad in 1:nrad
+        x = inputtglf_to_surrogate_input(all_inputs_tglf,irad)
+        y = [flux_solutions[isim][irad].ENERGY_FLUX_e for isim in 1:nsims]
+        ENERGY_FLUX_e = AbstractGPSurrogate(x,y, gp=GP(Matern32Kernel()), Σy=0.0001) 
+
+        y = [flux_solutions[isim][irad].ENERGY_FLUX_i for isim in 1:nsims]
+        ENERGY_FLUX_i = AbstractGPSurrogate(x,y, gp=GP(Matern32Kernel()), Σy=0.0001) 
+
+        y = [flux_solutions[isim][irad].PARTICLE_FLUX_e for isim in 1:nsims]
+        PARTICLE_FLUX_e = AbstractGPSurrogate(x,y, gp=GP(Matern32Kernel()), Σy=0.0001) 
+
+        PARTICLE_FLUX_i = []
+
+        y = [flux_solutions[isim][irad].STRESS_TOR_i for isim in 1:nsims]
+        STRESS_TOR_i = AbstractGPSurrogate(x,y, gp=GP(Matern32Kernel()), Σy=0.0001) 
+
+        push!(flux_surrogates, flux_surrogate(ENERGY_FLUX_e,ENERGY_FLUX_i,PARTICLE_FLUX_e,PARTICLE_FLUX_i,STRESS_TOR_i))
+    end
+    return flux_surrogates
+end function
+
+function inputtglf_to_surrogate_input(all_input_tglf,irad)
+    nsims = length(all_input_tglfs)
+    v = Array{Tuple{Vararg{Float64, 8}}}()
+
+    for isim in 1:nsims
+        push!(v,  (input_tglf.BETAE,
+        input_tglf.P_PRIME_LOC,
+        input_tglf.RLNS_1,
+        input_tglf.RLTS_1,
+        input_tglf.RLTS_2,
+        input_tglf.VEXB_SHEAR,
+        input_tglf.XNUE,
+        input_tglf.ZEFF))
+    end
+    return v
+
+end function
+
+"""
+    _step(actor::ActorSurrogate)
+
+Runs QLGYRO actor to evaluate the turbulence flux on a vector of gridpoints
+"""
+
+function _step(actor::ActorSurrogate)
+    par = actor.par
+    dd = actor.dd
+
+    if par.generate_surrogates
+         generate_surrogates(par.all_model_inputs, par.all_model_outputs)
+    end
+    if par.run_surrogates
+
+        actor.flux_solutions =  Vector{flux_solution}()
+        for irad in 1:nrad
+            all_input_tglf[end]
+            x = inputtglf_to_surrogate_input(par.all_model_inputs,irad)[end]
+            push!(actor.flux_solutions,flux_surrogates(x))
+        end
+    end
+    return actor
+end
+
+"""
+    _finalize(actor::ActorSurrogate)
+
+Writes results to dd.core_transport
+"""
+function _finalize(actor::ActorSurrogate)
+    dd = actor.dd
+    par = actor.par
+    cp1d = dd.core_profiles.profiles_1d[]
+    eqt = dd.equilibrium.time_slice[]
+
+    model = resize!(dd.core_transport.model, :anomalous; wipe=false)
+    model.identifier.name = string(par.model)
+    m1d = resize!(model.profiles_1d)
+    m1d.grid_flux.rho_tor_norm = par.rho_transport
+
+    IMAS.flux_gacode_to_fuse((:electron_energy_flux, :ion_energy_flux,  :electron_particle_flux, :ion_particle_flux, :momentum_flux), actor.flux_solutions, m1d, eqt, cp1d)
+
+    return actor
+end
+
+function Base.show(io::IO, ::MIME"text/plain", input::InputQLGYRO)
+    for field_name in fieldnames(typeof(input))
+        println(io, " $field_name = $(getfield(input,field_name))")
+    end
+end
+
+
+struct flux_surrogate{T<:Any}
+    ENERGY_FLUX_e::T
+    ENERGY_FLUX_i::T
+    PARTICLE_FLUX_e::T
+    PARTICLE_FLUX_i::Vector{T}
+    STRESS_TOR_i::T
+end