add ion_particle_flux and refactor NEO parsing

src/NEO.jl

@@ -83,50 +83,39 @@ function run_neo(input_neo::InputNEO)
     run(Cmd(`bash command.sh`; dir=folder))
 
     ### parse outputs ###
-    fluxes = Float64[]
-
-    tmp = open(joinpath(folder, "out.neo.transport_flux"), "r") do io
+    tmp_fluxes = Float64[]
+    open(joinpath(folder, "out.neo.transport_flux"), "r") do io
         for line in eachline(io)
             if !startswith(line, "#")
                 for word in split(line)
                     val = tryparse(Float64, word)
                     if val !== nothing
-                        push!(fluxes, val)
+                        push!(tmp_fluxes, val)
                     end
                 end
             end
         end
     end
-
     loc_first_tgyro = (4 * input_neo.N_SPECIES * 2) + 1
-    tgyro_fluxes = fluxes[loc_first_tgyro:end]
-
-    flux_solution = NEO.flux_solution()
-
-    for i in range(0, input_neo.N_SPECIES - 1)
-        species = i + 1
-        setfield!(flux_solution, Symbol("PARTICLE_FLUX_$species"), tgyro_fluxes[2+(i*4)])
-        setfield!(flux_solution, Symbol("ENERGY_FLUX_$species"), tgyro_fluxes[3+(i*4)])
-        setfield!(flux_solution, Symbol("MOMENTUM_FLUX_$species"), tgyro_fluxes[4+(i*4)])
-    end
-
-    energy_flux_electrons = getfield(flux_solution, Symbol("ENERGY_FLUX_$(input_neo.N_SPECIES)")) # electrons are always the last species in the list
-    total_ion_energy_flux = -energy_flux_electrons # exclude electron energy flux from total ion energy flux
-    particle_flux_electrons = getfield(flux_solution, Symbol("PARTICLE_FLUX_$(input_neo.N_SPECIES)"))
-
-    for field in fieldnames(NEO.flux_solution)
-        if ismissing(getfield(flux_solution, field))
-            setfield!(flux_solution, field, 0.0)
-        end
-
-        if occursin("ENERGY", string(field))
-            total_ion_energy_flux += getfield(flux_solution, field)
-        end
-    end
-
-    total_fluxes = IMAS.flux_solution(particle_flux_electrons, 0.0, energy_flux_electrons, total_ion_energy_flux)
-
-    return total_fluxes
+    tgyro_fluxes = tmp_fluxes[loc_first_tgyro:end]
+
+    # figure out indexes
+    e_index = [input_neo.N_SPECIES]
+    i_index = collect(1:input_neo.N_SPECIES-1)
+    particle_index(index) = 2 .+ ((index .- 1) .* 4)
+    energy_index(index) = 3 .+ ((index .- 1) .* 4)
+    momentum_index(index) = 4 .+ ((index .- 1) .* 4)
+
+    # sort fluxes
+    electrons_energy_flux = only(tgyro_fluxes[energy_index(e_index)])
+    electrons_particle_flux = only(tgyro_fluxes[particle_index(e_index)])
+    ion_particle_flux = tgyro_fluxes[particle_index(i_index)]
+    ion_total_energy_flux = sum(tgyro_fluxes[energy_index(i_index)])
+    ion_total_momentum_flux = sum(tgyro_fluxes[momentum_index(i_index)])
+
+    # assign fluxes to flux_solution structure
+    sol = IMAS.flux_solution(electrons_energy_flux, ion_total_energy_flux, electrons_particle_flux, ion_particle_flux, ion_total_momentum_flux)
+    return sol
 end
 
 const document = Dict()

src/chang_hinton.jl

@@ -6,6 +6,7 @@
         iion::Integer)
 
 Calculates the neoclassical flux using Chang-Hinton model which has been modified assuming Zi = 1, and ni=ne
+
 Ref: C.S. Chang, F.L. Hinton, Phys. Fluids 25, 1493–1494 (1982), https://doi.org/10.1063/1.863934
 """
 function changhinton(
@@ -103,6 +104,7 @@ function changhinton(
 
     qneo_gb = neo_rho_star_in^2
 
-    sol = IMAS.flux_solution(0.0, 0.0, 0.0, efluxi / qneo_gb)
+    # assign fluxes to flux_solution structure
+    sol = IMAS.flux_solution(0.0, efluxi / qneo_gb, 0.0, Float64[], 0.0)
     return sol
 end

src/hirshman_sigmar.jl

@@ -411,13 +411,14 @@ function HS_to_GB(HS_solution::Tuple{Vector{Float64},Vector{Float64}}, dd::IMAS.
     pflux_norm = pflux_multi ./ Gamma_neo_GB
     eflux_norm = eflux_multi ./ Q_neo_GB
 
-    particle_flux_e = pflux_norm[end]
     energy_flux_e = eflux_norm[end]
+    energy_flux_i = sum(eflux_norm[1:end-1])
+    particle_flux_e = pflux_norm[end]
+    particle_flux_i = pflux_norm[1:end-1]
 
-    energy_flux_i = sum(eflux_norm) - energy_flux_e
-
-    HS_fluxes = IMAS.flux_solution(particle_flux_e, 0.0, energy_flux_e, energy_flux_i)
-    return HS_fluxes
+    # assign fluxes to flux_solution structure
+    sol = IMAS.flux_solution(energy_flux_e, energy_flux_i, particle_flux_e, particle_flux_i, 0.0)
+    return sol
 end
 
 """