diff --git a/Project.toml b/Project.toml index de1e4d6..f99f312 100644 --- a/Project.toml +++ b/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" diff --git a/src/SD4SOLPS.jl b/src/SD4SOLPS.jl index 8c99eb1..4bcec72 100644 --- a/src/SD4SOLPS.jl +++ b/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 diff --git a/src/repair_eq.jl b/src/repair_eq.jl index 3f82e57..05af026 100644 --- a/src/repair_eq.jl +++ b/src/repair_eq.jl @@ -14,7 +14,7 @@ export check_rho_1d """ check_rho_1d( - dd::IMASDD.dd; + dd::IMAS.dd; time_slice::Int=1, throw_on_fail::Bool=false, )::Bool @@ -22,7 +22,7 @@ export check_rho_1d 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) diff --git a/src/supersize_profile.jl b/src/supersize_profile.jl index cd70054..ce424b8 100644 --- a/src/supersize_profile.jl +++ b/src/supersize_profile.jl @@ -2,7 +2,7 @@ Utilities for extrapolating profiles """ -using IMASDD: IMASDD +using IMAS: IMAS using Interpolations: Interpolations using GGDUtils: GGDUtils, get_grid_subset, add_subset_element!, get_subset_boundary, @@ -66,7 +66,7 @@ function extrapolate_core( edge_quantity::Vector{Float64}, rho_output::Vector{Float64}, )::Vector{Float64} - grad = IMASDD.gradient(edge_rho, edge_quantity) + grad = IMAS.gradient(edge_rho, edge_quantity) gf = grad[1] rf = edge_rho[1] gmid = -abs(gf) / 4.0 @@ -103,7 +103,7 @@ end """ fill_in_extrapolated_core_profile!( - dd::IMASDD.dd, + dd::IMAS.dd, quantity_name::String; method::String="simple", eq_time_idx::Int=1, @@ -143,7 +143,7 @@ Input arguments: this index should be set to -5. """ function fill_in_extrapolated_core_profile!( - dd::IMASDD.dd, + dd::IMAS.dd, quantity_name::String; method::String="simple", eq_time_idx::Int=1, @@ -243,7 +243,7 @@ function fill_in_extrapolated_core_profile!( ).(psi_for_quantity[in_bounds]) # Make sure the output 1D rho grid exists; create it if needed - if IMASDD.ismissing(dd.core_profiles.profiles_1d[it].grid, :rho_tor_norm) + if IMAS.ismissing(dd.core_profiles.profiles_1d[it].grid, :rho_tor_norm) # If you don't like this default, then you should write grid.rho_tor_norm # before calling this function. dd.core_profiles.profiles_1d[it].grid.rho_tor_norm = @@ -305,7 +305,7 @@ function extrapolate_edge_exp( end """ - prep_flux_map(dd::IMASDD.dd; eq_time_idx::Int=1, eq_profiles_2d_idx::Int=1) + prep_flux_map(dd::IMAS.dd; eq_time_idx::Int=1, eq_profiles_2d_idx::Int=1) Reads equilibrium data and extracts/derives some useful quantities. This is very basic, but it was being repeated and that's a no-no. @@ -316,7 +316,7 @@ Returns: - normalized poloidal flux on the equilibrium grid - a linear interpolation of norm pol flux vs. R and Z, ready to be evaluated """ -function prep_flux_map(dd::IMASDD.dd; eq_time_idx::Int=1, eq_profiles_2d_idx::Int=1) +function prep_flux_map(dd::IMAS.dd; eq_time_idx::Int=1, eq_profiles_2d_idx::Int=1) eqt = dd.equilibrium.time_slice[eq_time_idx] p2 = eqt.profiles_2d[eq_profiles_2d_idx] r_eq = p2.grid.dim1 @@ -331,7 +331,7 @@ end """ mesh_psi_spacing( - dd::IMASDD.dd; + dd::IMAS.dd; eq_time_idx::Int=1, eq_profiles_2d_idx::Int=1, grid_ggd_idx::Int=1, @@ -364,7 +364,7 @@ Input Arguments: same as the spacing at the edge of the mesh, or the same as the average spacing """ function mesh_psi_spacing( - dd::IMASDD.dd; + dd::IMAS.dd; eq_time_idx::Int=1, eq_profiles_2d_idx::Int=1, grid_ggd_idx::Int=1, @@ -427,7 +427,7 @@ end """ pick_extension_psi_range( - dd::IMASDD.dd; + dd::IMAS.dd; eq_time_idx::Int=1, eq_profiles_2d_idx::Int=1, grid_ggd_idx::Int=1, @@ -440,7 +440,7 @@ out to the most distant (in flux space) point on the limiting surface. Returns a vector of `psi_N` levels. """ function pick_extension_psi_range( - dd::IMASDD.dd; + dd::IMAS.dd; eq_time_idx::Int=1, eq_profiles_2d_idx::Int=1, grid_ggd_idx::Int=1, @@ -489,8 +489,8 @@ end """ pick_mesh_ext_starting_points( - grid_ggd::IMASDD.edge_profiles__grid_ggd, - space::IMASDD.edge_profiles__grid_ggd___space, + grid_ggd::IMAS.edge_profiles__grid_ggd, + space::IMAS.edge_profiles__grid_ggd___space, )::Tuple{Vector{Float64}, Vector{Float64}} Picks starting points for the radial lines of the mesh extension. The strategy @@ -499,8 +499,8 @@ gradient (of psi_N) to extend these gridlines outward. Returns a tuple with vectors of R and Z starting points. """ function pick_mesh_ext_starting_points( - grid_ggd::IMASDD.edge_profiles__grid_ggd, - space::IMASDD.edge_profiles__grid_ggd___space, + grid_ggd::IMAS.edge_profiles__grid_ggd, + space::IMAS.edge_profiles__grid_ggd___space, )::Tuple{Vector{Float64}, Vector{Float64}} # Choose starting points for the orthogonal (to the contour) gridlines # Use the existing cells of the standard mesh @@ -590,7 +590,7 @@ function mesh_ext_follow_grad( end # Step along the paths of steepest descent to populate the mesh. - dpsindr, dpsindz = IMASDD.gradient(r_eq, z_eq, psin_eq) + dpsindr, dpsindz = IMAS.gradient(r_eq, z_eq, psin_eq) dpdr = Interpolations.linear_interpolation((r_eq, z_eq), dpsindr) dpdz = Interpolations.linear_interpolation((r_eq, z_eq), dpsindz) rlim = (minimum(r_eq), maximum(r_eq)) @@ -625,7 +625,7 @@ end """ modify_mesh_ext_near_x!( - eqt::IMASDD.equilibrium__time_slice, + eqt::IMAS.equilibrium__time_slice, mesh_r::Matrix{Float64}, mesh_z::Matrix{Float64}, ) @@ -640,7 +640,7 @@ Input Arguments: - `mesh_z`: matrix of Z values for the extended mesh """ function modify_mesh_ext_near_x!( - eqt::IMASDD.equilibrium__time_slice, + eqt::IMAS.equilibrium__time_slice, mesh_r::Matrix{Float64}, mesh_z::Matrix{Float64}, ) @@ -723,8 +723,8 @@ end """ record_regular_mesh!( - grid_ggd::IMASDD.edge_profiles__grid_ggd, - space::IMASDD.edge_profiles__grid_ggd___space, + grid_ggd::IMAS.edge_profiles__grid_ggd, + space::IMAS.edge_profiles__grid_ggd___space, mesh_r::Matrix{Float64}, mesh_z::Matrix{Float64}, cut::Int, @@ -742,8 +742,8 @@ Records arrays of mesh data from regular 2D arrays into the DD and the next index. """ function record_regular_mesh!( - grid_ggd::IMASDD.edge_profiles__grid_ggd, - space::IMASDD.edge_profiles__grid_ggd___space, + grid_ggd::IMAS.edge_profiles__grid_ggd, + space::IMAS.edge_profiles__grid_ggd___space, mesh_r::Matrix{Float64}, mesh_z::Matrix{Float64}, cut::Int, @@ -899,7 +899,7 @@ end """ cached_mesh_extension!( - dd::IMASDD.dd, + dd::IMAS.dd, eqdsk_file::String, b2fgmtry::String; eq_time_idx::Int=1, @@ -926,7 +926,7 @@ Input Arguments: - `clear_cache`: delete any existing cache file (for use in testing) """ function cached_mesh_extension!( - dd::IMASDD.dd, + dd::IMAS.dd, eqdsk_file::String, b2fgmtry::String; eq_time_idx::Int=1, @@ -991,7 +991,7 @@ end """ mesh_extension_sol!( - dd::IMASDD.dd; + dd::IMAS.dd; eq_time_idx::Int=1, eq_profiles_2d_idx::Int=1, grid_ggd_idx::Int=1, @@ -1001,7 +1001,7 @@ end Extends the mesh out into the SOL """ function mesh_extension_sol!( - dd::IMASDD.dd; + dd::IMAS.dd; eq_time_idx::Int=1, eq_profiles_2d_idx::Int=1, grid_ggd_idx::Int=1, @@ -1048,7 +1048,7 @@ end """ fill_in_extrapolated_edge_profile!( - dd::IMASDD.dd, + dd::IMAS.dd, quantity_name::String; method::String="simple", eq_time_idx::Int=1, @@ -1057,7 +1057,7 @@ end JUST A PLACEHOLDER FOR NOW. DOESN'T ACTUALLY WORK YET. """ function fill_in_extrapolated_edge_profile!( - dd::IMASDD.dd, + dd::IMAS.dd, quantity_name::String; method::String="simple", eq_time_idx::Int=1, diff --git a/test/runtests.jl b/test/runtests.jl index 7f28771..235e3c9 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,6 +1,6 @@ using SD4SOLPS: SD4SOLPS using SOLPS2IMAS: SOLPS2IMAS -using IMASDD: IMASDD +using IMAS: IMAS using EFIT: EFIT using Plots using Test @@ -11,7 +11,7 @@ using GGDUtils: GGDUtils, get_grid_subset function parse_commandline() # Define newARGS = ["--yourflag"] to run only tests on your flags when including runtests.jl - localARGS = @isdefined(newARGS) ? newARGS : ARGS # Thanks https://stackoverflow.com/a/44978474/6605826 + localARGS = (@isdefined(newARGS) && newARGS !== nothing) ? newARGS : ARGS # Thanks https://stackoverflow.com/a/44978474/6605826 s = ArgParse.ArgParseSettings(; description="Run tests. Default is all tests.") ArgParse.add_arg_table!(s, @@ -270,7 +270,7 @@ if args["heavy_utilities"] # Test for sweeping 1D core profiles into 2D R,Z # (or anyway evaluating them at any R,Z location) - dd = IMASDD.dd() + dd = IMAS.dd() eqdsk_file = splitdir(pathof(SD4SOLPS))[1] * "/../sample/geqdsk_iter_small_sample" SD4SOLPS.geqdsk_to_imas!(eqdsk_file, dd) @@ -306,7 +306,7 @@ end if args["repair_eq"] @testset "repair_eq" begin # Prepare sample - dd = IMASDD.dd() + dd = IMAS.dd() eqdsk = splitdir(pathof(SD4SOLPS))[1] * "/../sample/geqdsk_iter_small_sample" SD4SOLPS.geqdsk_to_imas!(eqdsk, dd) # Make sure rho is missing @@ -338,7 +338,7 @@ if args["geqdsk_to_imas"] tslice = 1 for sample_file ∈ sample_files println(sample_file) - dd = IMASDD.dd() + dd = IMAS.dd() if endswith(sample_file, "00") eqdsk_time = parse(Float64, split(sample_file, ".")[end]) / 1000.0 else @@ -448,5 +448,7 @@ if args["preparation"] @test size(psirz) == (length(r), length(z)) println(out_file) @test isfile(out_file) + print("imas2json timing: ") + @time IMAS.imas2json(dd, filename * ".json", strict=true, freeze=false) end end