Merge pull request #121 from ProjectTorreyPines/add_sanitize

Add sanitize

Author: TimSlendebroek

src/expressions/dynamic.jl

@@ -118,14 +118,6 @@ dyexp["core_profiles.profiles_1d[:].j_tor"] =
         Jpar_2_Jtor(rho_tor_norm, profiles_1d.j_total, true, eqt)
     end
 
-#  core_profiles.vacuum_toroidal_field  #
-
-dyexp["core_profiles.vacuum_toroidal_field.b0"] =
-    (time; dd, _...) -> vacuum_r0_b0_time(dd, time)[2]
-
-dyexp["core_profiles.vacuum_toroidal_field.r0"] =
-    (; dd, _...) -> vacuum_r0_b0_time(dd)[1]
-
 #  core_profiles.global_quantities  #
 
 dyexp["core_profiles.global_quantities.current_non_inductive"] =
@@ -153,13 +145,6 @@ dyexp["core_profiles.profiles_1d[:].time"] =
 #= ============ =#
 # core_transport #
 #= ============ =#
-dyexp["core_transport.vacuum_toroidal_field.b0"] =
-    (time; dd, _...) -> vacuum_r0_b0_time(dd, time)[2]
-
-dyexp["core_transport.vacuum_toroidal_field.r0"] =
-    (; dd, _...) -> vacuum_r0_b0_time(dd)[1]
-
-
 dyexp["core_transport.model[:].profiles_1d[:].time"] =
     (; core_transport, profiles_1d_index, _...) -> begin
         return core_transport.time[profiles_1d_index]
@@ -168,12 +153,6 @@ dyexp["core_transport.model[:].profiles_1d[:].time"] =
 #= ========= =#
 # equilibrium #
 #= ========= =#
-# IMAS does not hold B0 information in a given time slice, but we can get that info from `B0=f/R0`
-# This trick propagates the B0 information to a time_slice even when that time_slice has not been initialized with profiles_1d data
-dyexp["equilibrium.time_slice[:].profiles_1d.f"] =
-    (psi; equilibrium, time_slice_index, _...) ->
-        (psi === missing ? [1] : ones(size(psi))) .* (equilibrium.vacuum_toroidal_field.b0[time_slice_index] * equilibrium.vacuum_toroidal_field.r0)
-
 dyexp["equilibrium.time_slice[:].global_quantities.energy_mhd"] =
     (; time_slice, _...) -> 3 / 2 * integrate(time_slice.profiles_1d.volume, time_slice.profiles_1d.pressure)
 
@@ -198,9 +177,15 @@ dyexp["equilibrium.time_slice[:].global_quantities.magnetic_axis.r"] =
 dyexp["equilibrium.time_slice[:].global_quantities.magnetic_axis.z"] =
     (; time_slice, _...) -> time_slice.profiles_1d.geometric_axis.z[1]
 
+
+dyexp["equilibrium.time_slice[:].global_quantities.vacuum_toroidal_field.b0"] =
+    (; dd, time_slice, _...) -> get_time_array(dd.pulse_schedule.tf.b_field_tor_vacuum, :reference, time_slice.time, :linear)
+
+dyexp["equilibrium.time_slice[:].global_quantities.vacuum_toroidal_field.r0"] =
+    (; dd, _...) -> dd.pulse_schedule.tf.r0
+
 dyexp["equilibrium.time_slice[:].global_quantities.magnetic_axis.b_field_tor"] =
-    (; equilibrium, time_slice_index, _...) ->
-        equilibrium.vacuum_toroidal_field.b0[time_slice_index] * equilibrium.vacuum_toroidal_field.r0 / equilibrium.time_slice[time_slice_index].boundary.geometric_axis.r
+    (; time_slice, _...) -> time_slice.global_quantities.vacuum_toroidal_field.b0 * time_slice.global_quantities.vacuum_toroidal_field.r0 / time_slice.boundary.geometric_axis.r
 
 dyexp["equilibrium.time_slice[:].profiles_1d.geometric_axis.r"] =
     (psi; time_slice, _...) -> (time_slice.profiles_1d.r_outboard .+ time_slice.profiles_1d.r_inboard) .* 0.5
@@ -209,10 +194,25 @@ dyexp["equilibrium.time_slice[:].profiles_1d.geometric_axis.z"] =
     (psi; time_slice, _...) -> psi .* 0.0 .+ time_slice.global_quantities.magnetic_axis.z
 
 dyexp["equilibrium.time_slice[:].boundary.geometric_axis.r"] =
-    (; time_slice, _...) -> time_slice.profiles_1d.geometric_axis.r[end]
+    (; time_slice, _...) -> begin
+        if !ismissing(time_slice.profiles_1d.geometric_axis, :r)
+            return time_slice.profiles_1d.geometric_axis.r[end]
+        else
+            minR, maxR = extrema(time_slice.boundary.outline.r)
+            return (minR + maxR) / 2
+        end
+    end
 
 dyexp["equilibrium.time_slice[:].boundary.geometric_axis.z"] =
-    (; time_slice, _...) -> time_slice.profiles_1d.geometric_axis.z[end]
+    (; time_slice, _...) -> begin
+        if !ismissing(time_slice.profiles_1d.geometric_axis, :z)
+            return time_slice.profiles_1d.geometric_axis.z[end]
+        else
+            minZ, maxZ = extrema(time_slice.boundary.outline.z)
+            return (minZ + maxZ) / 2
+        end
+    end
+
 
 dyexp["equilibrium.time_slice[:].boundary.minor_radius"] =
     (; time_slice, _...) -> (time_slice.profiles_1d.r_outboard[end] - time_slice.profiles_1d.r_inboard[end]) * 0.5
@@ -287,13 +287,6 @@ dyexp["equilibrium.time_slice[:].profiles_1d.dpsi_drho_tor"] =
 dyexp["equilibrium.time_slice[:].profiles_1d.psi_norm"] =
     (psi; _...) -> norm01(psi)
 
-
-dyexp["equilibrium.vacuum_toroidal_field.b0"] =
-    (time; dd, _...) -> vacuum_r0_b0_time(dd, time)[2]
-
-dyexp["equilibrium.vacuum_toroidal_field.r0"] =
-    (; dd, _...) -> vacuum_r0_b0_time(dd)[1]
-
 # 2D
 dyexp["equilibrium.time_slice[:].profiles_2d[:].r"] =
     (dim1, dim2; _...) -> ones(length(dim2))' .* dim1
@@ -412,14 +405,6 @@ dyexp["core_sources.source[:].profiles_1d[:].ion[:].particles"] =
         gradient(profiles_1d.grid.volume, ion.particles_inside)
     end
 
-
-dyexp["core_sources.vacuum_toroidal_field.b0"] =
-    (time; dd, _...) -> vacuum_r0_b0_time(dd, time)[2]
-
-dyexp["core_sources.vacuum_toroidal_field.r0"] =
-    (; dd, _...) -> vacuum_r0_b0_time(dd)[1]
-
-
 dyexp["core_sources.source[:].profiles_1d[:].time"] =
     (; core_sources, profiles_1d_index, _...) -> begin
         return core_sources.time[profiles_1d_index]
@@ -587,6 +572,18 @@ dyexp["pulse_schedule.time"] =
         return sort!(unique(all_times))
     end
 
+dyexp["pulse_schedule.tf.b_field_tor_vacuum_r.reference"] =
+    (time; tf, _...) ->  tf.r0 .* tf.b_field_tor_vacuum.reference 
+
+dyexp["pulse_schedule.tf.r0"] =
+    (; dd, _...) ->  dd.equilibrium.vacuum_toroidal_field.r0
+
+dyexp["pulse_schedule.tf.b_field_tor_vacuum.reference"] =
+    (time; dd, _...) ->  dd.equilibrium.vacuum_toroidal_field.b0
+
+dyexp["pulse_schedule.tf.time"] =
+    (time; dd, _...) ->  dd.equilibrium.time
+
 #= ========= =#
 #  stability  #
 #= ========= =#
@@ -617,12 +614,6 @@ dyexp["summary.fusion.power.value"] = # NOTE: This is the fusion power that is c
 dyexp["summary.global_quantities.ip.value"] =
     (time; dd, summary, _...) -> [dd.equilibrium.time_slice[Float64(time0)].global_quantities.ip for time0 in time]
 
-dyexp["summary.global_quantities.b0.value"] =
-    (time; dd, _...) -> vacuum_r0_b0_time(dd, time)[2]
-
-dyexp["summary.global_quantities.r0.value"] =
-    (; dd, summary, _...) -> vacuum_r0_b0_time(dd)[1]
-
 dyexp["summary.global_quantities.current_bootstrap.value"] =
     (time; dd, summary, _...) -> begin
         tmp = eltype(summary)[]

src/extract.jl

@@ -34,7 +34,7 @@ function update_ExtractFunctionsLibrary!()
     ExtractLibFunction(:geometry, :Volume, "m³", dd -> dd.equilibrium.time_slice[].profiles_1d.volume[end])
     ExtractLibFunction(:geometry, :Surface, "m²", dd -> dd.equilibrium.time_slice[].profiles_1d.surface[end])
 
-    ExtractLibFunction(:equilibrium, :B0, "T", dd -> @ddtime(dd.summary.global_quantities.b0.value))
+    ExtractLibFunction(:equilibrium, :B0, "T", dd -> dd.equilibrium.time_slice[].global_quantities.vacuum_toroidal_field.b0)
     ExtractLibFunction(:equilibrium, :ip, "MA", dd -> @ddtime(dd.summary.global_quantities.ip.value) / 1e6)
     ExtractLibFunction(:equilibrium, :q95, "-", dd -> dd.equilibrium.time_slice[].global_quantities.q_95)
     ExtractLibFunction(:equilibrium, Symbol("<Bpol>"), "T", dd -> IMAS.Bpol(EFL[:a](dd), EFL[:κ](dd), EFL[:ip](dd) * 1e6))

src/get_from.jl

@@ -17,6 +17,18 @@ function get_from(dd::IMAS.dd{T}, what::Type{Val{:ip}}, from_where::Symbol; time
     end
 end
 
+# vacuum_r0_b0 [m], [T]
+function get_from(dd::IMAS.dd, what::Type{Val{:vacuum_r0_b0}}, from_where::Symbol; time0::Float64=dd.global_time)
+    if from_where == :equilibrium
+        eqt = dd.equilibrium.time_slice[time0]
+        return (r0=eqt.global_quantities.vacuum_toroidal_field.r0, b0=eqt.global_quantities.vacuum_toroidal_field.b0)
+    elseif from_where == :pulse_schedule
+        return (r0=dd.pulse_schedule.tf.r0, b0=IMAS.get_time_array(dd.pulse_schedule.tf.b_field_tor_vacuum, :reference, time0, :linear))
+    else
+        error("`get_from(dd, $what, $from_where)` doesn't exist yet")
+    end
+end
+
 # vloop [V]
 function get_from(dd::IMAS.dd{T}, what::Type{Val{:vloop}}, from_where::Symbol; time0::Float64=dd.global_time)::T where {T<:Real}
     if from_where == :equilibrium

src/physics/boundary.jl

@@ -256,16 +256,17 @@ function boundary_shape(;
     z = vcat(z1, z2, z3, z4)
     zref = vcat(z1ref, z2ref, z3ref, z4ref)
 
-    return r, z, zref
+    return (r=r, z=z, zref=zref)
 end
 
-function boundary(pc::IMAS.pulse_schedule__position_control, time0::Float64)
-    return [extrap1d(interp1d_itp(pc.time, pcb.r.reference); first=:flat, last=:flat).(time0) for pcb in pc.boundary_outline],
-    [extrap1d(interp1d_itp(pc.time, pcb.z.reference); first=:flat, last=:flat).(time0) for pcb in pc.boundary_outline]
+function boundary(pc::IMAS.pulse_schedule__position_control{T}, time0::Float64) where {T<:Real}
+    return (r=T[extrap1d(interp1d_itp(pc.time, pcb.r.reference); first=:flat, last=:flat).(time0) for pcb in pc.boundary_outline],
+        z=T[extrap1d(interp1d_itp(pc.time, pcb.z.reference); first=:flat, last=:flat).(time0) for pcb in pc.boundary_outline])
 end
 
-function boundary(pc::IMAS.pulse_schedule__position_control, time_index::Int)
-    return [pcb.r.reference[time_index] for pcb in pc.boundary_outline], [pcb.z.reference[time_index] for pcb in pc.boundary_outline]
+function boundary(pc::IMAS.pulse_schedule__position_control{T}, time_index::Int) where {T<:Real}
+    return (r=T[pcb.r.reference[time_index] for pcb in pc.boundary_outline],
+        z=T[pcb.z.reference[time_index] for pcb in pc.boundary_outline])
 end
 
 """

src/physics/currents.jl

@@ -142,7 +142,7 @@ function vloop(cp1d::IMAS.core_profiles__profiles_1d{T}, eqt::IMAS.equilibrium__
     rho_eq = eqt.profiles_1d.rho_tor_norm
     F = IMAS.interp1d(rho_eq, eqt.profiles_1d.f, :cubic).(rho_tor_norm)
     gm1 = IMAS.interp1d(rho_eq, eqt.profiles_1d.gm1, :cubic).(rho_tor_norm) # <R⁻²>
-    _, B0 = IMAS.vacuum_r0_b0(eqt)
+    _, B0 = eqt.global_quantities.vacuum_toroidal_field.r0, eqt.global_quantities.vacuum_toroidal_field.b0
     Vls = 2π .* cp1d.j_ohmic .* B0 ./ (cp1d.conductivity_parallel .* F .* gm1)
     if method === :area
         return integrate(cp1d.grid.area, Vls) / cp1d.grid.area[end]

src/physics/equilibrium.jl

@@ -123,119 +123,13 @@ function symmetrize_equilibrium!(eqt::IMAS.equilibrium__time_slice)
     tweak_psi_to_match_psilcfs!(eqt)
 end
 
-"""
-    vacuum_r0_b0(eqt::IMAS.equilibrium__time_slice{T})::Tuple{T,T} where {T<:Real}
-
-Returns vacuum R0 and B0 of a given equilibrium time slice
-"""
-function vacuum_r0_b0(eqt::IMAS.equilibrium__time_slice{T})::Tuple{T,T} where {T<:Real}
-    eq = top_ids(eqt)
-    if eq !== nothing && !ismissing(eq.vacuum_toroidal_field, :r0) && !ismissing(eq.vacuum_toroidal_field, :b0)
-        R0 = eq.vacuum_toroidal_field.r0
-        B0 = get_time_array(eq.vacuum_toroidal_field, :b0, eqt.time)
-    else
-        R0 = eqt.boundary.geometric_axis.r
-        B0 = eqt.profiles_1d.f[end] / R0
-    end
-    return R0, B0
-end
-
-"""
-    vacuum_r0_b0_time(dd::IMAS.dd{T}) where {T<:Real}
-
-Returns `R0` as well as the `B0` and `time` arrays.
-
-This function solves the issue that in IMAS the information about R0 and B0 is replicated across different IDSs.
-"""
-function vacuum_r0_b0_time(dd::IMAS.dd{T}) where {T<:Real}
-    source = Set{Symbol}()
-
-    idss_with_vacuum_toroidal_field = keys(dd)
-
-    # R0
-    if hasfield(typeof(dd), :tf) && hasdata(dd.tf, :r0)
-        R0 = dd.tf.r0::T
-        push!(source, :tf)
-    else
-        for name in idss_with_vacuum_toroidal_field
-            ids = getfield(dd, name)
-            if hasfield(typeof(ids), :vacuum_toroidal_field)
-                if hasdata(ids.vacuum_toroidal_field, :r0)
-                    R0 = ids.vacuum_toroidal_field.r0::T
-                    push!(source, name)
-                    break
-                end
-            end
-        end
-    end
-
-    # B0 and time vectors
-    # from: tf
-    if hasfield(typeof(dd), :tf) && hasdata(dd.tf.b_field_tor_vacuum_r, :data)
-        B0 = dd.tf.b_field_tor_vacuum_r.data::Vector{T} / R0
-        time = dd.tf.b_field_tor_vacuum_r.time::Vector{Float64}
-        push!(source, :tf)
-
-        # from: pulse_schedule
-    elseif hasdata(dd.pulse_schedule.tf.b_field_tor_vacuum_r, :reference)
-        B0 = dd.pulse_schedule.tf.b_field_tor_vacuum_r.reference::Vector{T} / R0
-        time = dd.pulse_schedule.tf.time::Vector{Float64}
-        push!(source, :pulse_schedule)
-
-        # from: all other IDSs that have that info
-    else
-        B0 = eltype(dd)[]
-        time = Float64[]
-        for name in idss_with_vacuum_toroidal_field
-            ids = getfield(dd, name)
-            if hasfield(typeof(ids), :vacuum_toroidal_field)
-                if hasdata(ids.vacuum_toroidal_field, :b0) && hasdata(ids, :time)
-                    B0_ = ids.vacuum_toroidal_field.b0::Vector{T}
-                    time_ = ids.time::Vector{Float64}
-                    append!(time, time_)
-                    append!(B0, B0_)
-                    reps = length(time_) - length(B0_)
-                    append!(B0, (B0_[end] for k in 1:reps))
-                    push!(source, name)
-                end
-            end
-        end
-        index = unique_indices(time)
-        B0 = B0[index]
-        time = time[index]
-    end
-
-    return R0, B0, time, source
-end
-
-"""
-    vacuum_r0_b0_time(dd::IMAS.dd, time::Vector{Float64})
-
-Returns R0 and B0 interpolated at a given set of times
-"""
-function vacuum_r0_b0_time(dd::IMAS.dd, time::Vector{Float64})
-    R0, B0_, time_, source = vacuum_r0_b0_time(dd)
-    B0 = extrap1d(interp1d_itp(time_, B0_); first=:flat, last=:flat).(time)
-    return R0, B0
-end
-
-"""
-    vacuum_r0_b0_time(dd::IMAS.dd{T}, time0::Float64)::Tuple{T,T}
-
-Returns R0 and B0 interpolated at one time
-"""
-function vacuum_r0_b0_time(dd::IMAS.dd{T}, time0::Float64)::Tuple{T,T} where{T<:Real}
-    R0, B0 = vacuum_r0_b0_time(dd, [time0])
-    return R0, B0[1]
-end
-
 """
     B0_geo(eqt::IMAS.equilibrium__time_slice{T})::T where{T<:Real}
 
 Returns vacuum B0 at the plasma geometric center
 """
 function B0_geo(eqt::IMAS.equilibrium__time_slice{T})::T where{T<:Real}
-    R0, B0 = vacuum_r0_b0(eqt)
+    R0, B0 = eqt.global_quantities.vacuum_toroidal_field.r0, eqt.global_quantities.vacuum_toroidal_field.b0
     Rgeo = eqt.boundary.geometric_axis.r
     return R0 .* B0 ./ Rgeo
 end

src/physics/fluxsurfaces.jl

@@ -649,7 +649,7 @@ Update flux surface averaged and geometric quantities for a given equilibrum IDS
 The original psi grid can be upsampled by a `upsample_factor` to get higher resolution flux surfaces
 """
 function flux_surfaces(eqt::equilibrium__time_slice; upsample_factor::Int=1)
-    R0, B0 = vacuum_r0_b0(eqt)
+    R0, B0 = eqt.global_quantities.vacuum_toroidal_field.r0, eqt.global_quantities.vacuum_toroidal_field.b0
     return flux_surfaces(eqt, B0, R0; upsample_factor)
 end
 

src/physics/pf_active.jl

@@ -52,7 +52,7 @@ function set_coils_function(coils::IDSvector{<:IMAS.pf_active__coil})
     geometry_types = name_2_index(coils[1].element[1].geometry)
     for coil in coils
         for element in coil.element
-            for geometry_type in keys(element.geometry)
+            for geometry_type in (:outline, :rectangle, :oblique)
                 if geometry_type != :geometry_type && !ismissing(getproperty(element.geometry, geometry_type), :r)
                     element.geometry.geometry_type = geometry_types[geometry_type]
                     break

src/physics/profiles.jl

@@ -220,7 +220,7 @@ function tau_e_h98(eqt::IMAS.equilibrium__time_slice, cp1d::IMAS.core_profiles__
         integrate(cp1d.grid.volume, sum(ion.density .* ion.element[1].a for ion in cp1d.ion if ion.element[1].z_n == 1.0)) /
         integrate(cp1d.grid.volume, sum(ion.density for ion in cp1d.ion if ion.element[1].z_n == 1.0))
 
-    R0, B0 = vacuum_r0_b0(eqt)
+    R0, B0 = eqt.global_quantities.vacuum_toroidal_field.r0, eqt.global_quantities.vacuum_toroidal_field.b0
 
     tau98 = (
         0.0562 *
@@ -255,7 +255,7 @@ function tau_e_ds03(eqt::IMAS.equilibrium__time_slice, cp1d::IMAS.core_profiles_
         integrate(cp1d.grid.volume, sum(ion.density .* ion.element[1].a for ion in cp1d.ion if ion.element[1].z_n == 1.0)) /
         integrate(cp1d.grid.volume, sum(ion.density for ion in cp1d.ion if ion.element[1].z_n == 1.0))
 
-    R0, B0 = vacuum_r0_b0(eqt)
+    R0, B0 = eqt.global_quantities.vacuum_toroidal_field.r0, eqt.global_quantities.vacuum_toroidal_field.b0
 
     tauds03 = (
         0.028 *

src/physics/sol.jl

@@ -133,7 +133,7 @@ function sol(eqt::IMAS.equilibrium__time_slice, wall_r::Vector{T}, wall_z::Vecto
     end
 
     ############ 
-    R0, B0 = vacuum_r0_b0(eqt)
+    R0, B0 = eqt.global_quantities.vacuum_toroidal_field.r0, eqt.global_quantities.vacuum_toroidal_field.b0
     RA = eqt.global_quantities.magnetic_axis.r
     ZA = eqt.global_quantities.magnetic_axis.z
 
@@ -541,7 +541,7 @@ Parallel heat flux [W/m^2] at the outer midplane based on Eigh λ_q
 """
 function q_par_omp_eich(eqt::IMAS.equilibrium__time_slice, cp1d::IMAS.core_profiles__profiles_1d, core_sources::IMAS.core_sources)
     eq1d = eqt.profiles_1d
-    R0, B0 = vacuum_r0_b0(eqt)
+    R0, B0 = eqt.global_quantities.vacuum_toroidal_field.r0, eqt.global_quantities.vacuum_toroidal_field.b0
     R_omp = eq1d.r_outboard[end]
     Bt_omp = B0 * R0 / R_omp
     return q_pol_omp_eich(eqt, cp1d, core_sources) / sin(atan(Bpol_omp(eqt) / Bt_omp))