Switch IMASDD to IMAS and make adaptations

Project.toml

@@ -9,7 +9,7 @@ Contour = "d38c429a-6771-53c6-b99e-75d170b6e991"
 EFIT = "cda752c5-6b03-55a3-9e33-132a441b0c17"
 Fortran90Namelists = "8fb689aa-71ff-4044-8071-0cffc910b57d"
 GGDUtils = "b7b5e640-9b39-4803-84eb-376048795def"
-IMASDD = "06b86afa-9f21-11ec-2ef8-e51b8960cfc5"
+IMAS = "13ead8c1-b7d1-41bb-a6d0-5b8b65ed587a"
 Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
 PhysicalConstants = "5ad8b20f-a522-5ce9-bfc9-ddf1d5bda6ab"
@@ -28,14 +28,14 @@ Contour = "0.6.3"
 EFIT = "1.0.0"
 Fortran90Namelists = "0.1.1"
 GGDUtils = "1.0.2"
-IMASDD = "0.1, 1"
+IMAS = "1.2.2"
 Interpolations = "0.15.1"
 JSON = "0.21.4"
 PhysicalConstants = "0.2.3"
 PlotUtils = "1.4.1"
 Plots = "1.40.3"
 PolygonOps = "0.1.2"
-SOLPS2IMAS = "1.0.2"
+SOLPS2IMAS = "1.0.3"
 Statistics = "1.9.0"
 Unitful = "1.20.0"
 YAML = "0.4.11"

src/SD4SOLPS.jl

@@ -1,6 +1,6 @@
 module SD4SOLPS
 
-using IMASDD: IMASDD
+using IMAS: IMAS
 using SOLPS2IMAS: SOLPS2IMAS
 using EFIT: EFIT
 using Interpolations: Interpolations
@@ -67,7 +67,7 @@ end
 """
     geqdsk_to_imas!(
         eqdsk_file::String,
-        dd::IMASDD.dd;
+        dd::IMAS.dd;
         set_time::Union{Nothing, Float64}=nothing,
         time_index::Int=1,
     )
@@ -77,16 +77,17 @@ the IMAS DD structure.
 """
 function geqdsk_to_imas!(
     eqdsk_file::String,
-    dd::IMASDD.dd;
+    dd::IMAS.dd;
     set_time::Union{Nothing, Float64}=nothing,
     time_index::Int=1,
+    allow_boundary_flux_correction::Bool=false,
 )
     # https://github.com/JuliaFusion/EFIT.jl/blob/master/src/io.jl
     g = EFIT.readg(eqdsk_file; set_time=set_time)
     gfilename = split(eqdsk_file, "/")[end]
     # Copying ideas from OMFIT: omfit/omfit_classes/omfit_eqdsk.py / to_omas()
     eq = dd.equilibrium
-    if IMASDD.ismissing(eq, :time)
+    if IMAS.ismissing(eq, :time)
         eq.time = Array{Float64}(undef, time_index)
     end
     eq.time[time_index] = g.time
@@ -110,7 +111,7 @@ function geqdsk_to_imas!(
     b0[time_index] = g.bcentr
     eq.vacuum_toroidal_field.b0 = b0
 
-    if IMASDD.ismissing(dd.summary, :time)
+    if IMAS.ismissing(dd.summary, :time)
         dd.summary.time = Array{Float64}(undef, time_index)
     end
     dd.summary.time[time_index] = g.time
@@ -121,31 +122,56 @@ function geqdsk_to_imas!(
     dd.summary.global_quantities.b0.value = b0
     summarize = ["ip", "r0", "b0"]
 
+    # 2D
+    resize!(eqt.profiles_2d, 1)
+    p2 = eqt.profiles_2d[1]
+    p2.grid.dim1 = collect(g.r)
+    p2.grid.dim2 = collect(g.z)
+    p2.psi = g.psirz  # Not sure if transpose is correct (I have been getting away with this for some time and suspect it's okay)
+    p2.grid_type.index = 1  # 1 = rectangular, such as dim1 = R, dim2 = Z
+    p2.grid_type.name = "R-Z grid for flux map"
+    p2.grid_type.description = (
+        "A rectangular grid of points in R,Z on which poloidal " *
+        "magnetic flux psi is defined. The grid's dim1 is R, dim2 is Z."
+    )
+    # missing j_tor = pcurrt
+
+    if allow_boundary_flux_correction
+        # Check / correct simag. Intended for the case where simag is quoted imprecisely and the contour doesn't close
+        level = gq.psi_boundary + 0
+        paths, level = IMAS.flux_surface(eqt, level, :closed)
+        count = 0  # prevents inf loop if something goes wrong
+        while (length(paths) >= 1) && (count < 50)  # push boundary flux out until there is no closed boundary
+            level += (gq.psi_boundary - gq.psi_axis) * 0.001
+            paths, level = IMAS.flux_surface(eqt, level, :closed)
+            count += 1
+        end
+        println(" --- ")
+        while (length(paths) < 1) && (count < 200)  # pull boundary flux back in until there is a closed boundary
+            level -= (gq.psi_boundary - gq.psi_axis) * 0.0001
+            paths, level = IMAS.flux_surface(eqt, level, :closed)
+            count += 1
+        end
+        println(
+            "Correcting psi_boundary from $(gq.psi_boundary) to $level so contouring will find a closed flux surface.",
+        )
+        gq.psi_boundary = level
+    end
+
     # 1D
     p1 = eqt.profiles_1d
-    nprof = length(g.pres)
-    psi = collect(LinRange(g.simag, g.sibry, nprof))
-    p1.psi = psi
+    p1.psi = collect(g.psi)
     p1.f = g.fpol
     p1.pressure = g.pres
     p1.f_df_dpsi = g.ffprim
     p1.dpressure_dpsi = g.pprime
     p1.q = g.qpsi
     if hasproperty(g, :rhovn)
-        # rhovn is not in the original EFIT.jl but is added on a branch
         p1.rho_tor_norm = g.rhovn
     end
 
-    # 2D
-    resize!(eqt.profiles_2d, 1)
-    p2 = eqt.profiles_2d[1]
-    p2.grid.dim1 = collect(g.r)
-    p2.grid.dim2 = collect(g.z)
-    p2.psi = g.psirz  # Not sure if transpose is correct
-    # missing j_tor = pcurrt
-
     # Derived
-    psin1d = (psi .- g.simag) ./ (g.sibry - g.simag)
+    psin1d = (g.psi .- gq.psi_axis) ./ (gq.psi_boundary - gq.psi_axis)
     gq.magnetic_axis.b_field_tor = g.bcentr * g.rcentr / g.rmaxis
     gq.q_axis = g.qpsi[1]
     gq.q_95 = Interpolations.linear_interpolation(psin1d, g.qpsi)(0.95)
@@ -222,7 +248,8 @@ end
         output_format::String="json",
         eqdsk_set_time::Union{Nothing, Float64}=nothing,
         eq_time_index::Int=1,
-    )::IMASDD.dd
+        allow_boundary_flux_correction::Bool=false,
+    )::IMAS.dd
 
 Gathers SOLPS and EFIT files and loads them into IMAS structure. Extrapolates
 profiles as needed to get a complete picture.
@@ -235,7 +262,8 @@ function preparation(
     output_format::String="json",
     eqdsk_set_time::Union{Nothing, Float64}=nothing,
     eq_time_index::Int=1,
-)::IMASDD.dd
+    allow_boundary_flux_correction::Bool=false,
+)::IMAS.dd
     b2fgmtry, b2time, b2mn, eqdsk =
         find_files_in_allowed_folders(dirs...; eqdsk_file=eqdsk_file)
     println("Found source files:")
@@ -244,12 +272,40 @@ function preparation(
     println("    b2mn.dat = ", b2mn)
     println("    eqdsk = ", eqdsk)
 
-    dd = SOLPS2IMAS.solps2imas(b2fgmtry, b2time; b2mn=b2mn)
-    geqdsk_to_imas!(eqdsk, dd; set_time=eqdsk_set_time, time_index=eq_time_index)
-    # Repairs
+    dd = IMAS.dd()
+    geqdsk_to_imas!(
+        eqdsk,
+        dd;
+        set_time=eqdsk_set_time,
+        time_index=eq_time_index,
+        allow_boundary_flux_correction=allow_boundary_flux_correction,
+    )
+    # Fill out more equilibrium data
     add_rho_to_equilibrium!(dd)  # Doesn't do anything if rho is valid
+    dd.global_time = dd.equilibrium.time_slice[1].time
+    for eqt ∈ dd.equilibrium.time_slice
+        IMAS.flux_surfaces(eqt)
+    end
+    # Add SOLPS data
+    dd = SOLPS2IMAS.solps2imas(b2fgmtry, b2time; b2mn=b2mn, ids=dd)
     println("Loaded input data into IMAS DD")
 
+    # Core profiles
+    # Set timing
+    nt = length(dd.edge_profiles.ggd)
+    if length(dd.core_profiles.profiles_1d) < nt
+        resize!(dd.core_profiles.profiles_1d, nt)
+    end
+    if ismissing(dd.core_profiles, :time)
+        dd.core_profiles.time = Array{Float64}(undef, nt)
+    elseif length(dd.core_profiles.time) < nt
+        resize!(dd.core_profiles.time, nt)
+    end
+    for it ∈ 1:nt
+        dd.core_profiles.time[it] =
+            dd.core_profiles.profiles_1d[it].time = dd.edge_profiles.ggd[it].time
+    end
+    # Extrapolate profiles
     core_profiles = ["electrons.density", "electrons.temperature"]
     extrapolated_core_profiles = []
     for core_profile ∈ core_profiles
@@ -279,8 +335,10 @@ function preparation(
     # ... more profiles here
     println("Extrapolated edge profiles (but not really (placeholder only))")
 
+    print("Exporting to file: ")
     if output_format == "json"
-        IMASDD.imas2json(dd, filename * ".json")
+        println(filename * ".json")
+        IMAS.imas2json(dd, filename * ".json"; strict=true, freeze=false)
     else
         throw(ArgumentError(string("Unrecognized output format: ", output_format)))
     end

src/repair_eq.jl

@@ -14,15 +14,15 @@ export check_rho_1d
 
 """
     check_rho_1d(
-        dd::IMASDD.dd;
+        dd::IMAS.dd;
         time_slice::Int=1,
         throw_on_fail::Bool=false,
     )::Bool
 
 Checks to see if rho exists and is valid in the equilibrium 1d profiles
 """
 function check_rho_1d(
-    dd::IMASDD.dd;
+    dd::IMAS.dd;
     time_slice::Int=1,
     throw_on_fail::Bool=false,
 )::Bool
@@ -56,11 +56,11 @@ function check_rho_1d(
 end
 
 """
-    function add_rho_to_equilibrium(dd:IMASDD.dd)
+    function add_rho_to_equilibrium(dd:IMAS.dd)
 
 Adds equilibrium rho profile to the DD
 """
-function add_rho_to_equilibrium!(dd::IMASDD.dd)
+function add_rho_to_equilibrium!(dd::IMAS.dd)
     nt = length(dd.equilibrium.time_slice)
     if nt < 1
         println("No equilibrium time slices to work with; can't add rho")
@@ -80,10 +80,10 @@ function add_rho_to_equilibrium!(dd::IMASDD.dd)
             end
         end
         if (
-            if IMASDD.ismissing(eqt.profiles_1d, :phi)
+            if IMAS.ismissing(eqt.profiles_1d, :phi)
                 true
             else
-                IMASDD.isempty(eqt.profiles_1d.phi)
+                IMAS.isempty(eqt.profiles_1d.phi)
             end
         )
             eqt.profiles_1d.phi = Array{Float64}(undef, n)