Try to use IMAS time handling instead of passing around indices

JuliaFusion · Aug 8, 2024 · 314d04b · 314d04b
1 parent 1045e2d
commit 314d04b
Show file tree

Hide file tree

Showing 2 changed files with 105 additions and 105 deletions.
diff --git a/src/io.jl b/src/io.jl
@@ -262,6 +262,7 @@ function geqdsk2imas!(
     cocos_clockwise_phi::Bool=false,
     add_derived::Bool=false,
 )
+    # Should work if dd is empty
     if ismissing(dd.dataset_description.data_entry, :pulse)
         dd.dataset_description.data_entry.pulse = parse(Int, split(split(gs[1].file, ".")[end-1], "g")[2])
     end
@@ -302,17 +303,32 @@ function geqdsk2imas!(
     cocos_clockwise_phi::Bool=false,
     add_derived::Bool=false,
 )
+    if ismissing(dd.dataset_description.data_entry, :pulse)
+        dd.dataset_description.data_entry.pulse = parse(Int, split(split(gs[1].file, ".")[end-1], "g")[2])
+    end
     if wall !== nothing
         geqdsk2wall!(gs[1], wall; geqdsk_cocos=geqdsk_cocos, dd_cocos=dd_cocos, cocos_clockwise_phi=cocos_clockwise_phi)
     end
     nt = length(gs)
-    if length(eq.time_slice) < nt
+    if length(eq.time_slice) == 0  # Blank is okay, but if populated, it had better be compatible
         resize!(eq.time_slice, nt)
+        eq.time = [g.time for g in gs]
+    end
+
+    if ismissing(eq.vacuum_toroidal_field, :b0)
+        eq.vacuum_toroidal_field.b0 = Array{Float64}(undef, nt)
     end
-    for it in 1:nt
-        g = gs[it]
+    if length(eq.vacuum_toroidal_field.b0) == 0
+        resize!(eq.vacuum_toroidal_field.b0, nt)
+    end
+    # Write top level data
+    eq.time = [g.time for g in gs]
+    eq.vacuum_toroidal_field.b0 = [g.bcentr .* tc["B"] for g in gs]
+    eq.vacuum_toroidal_field.r0 = gs[1].rcentr
+
+    for g in gs
         geqdsk2imas!(
-            g, eq, it;
+            g, eq;
             geqdsk_cocos=geqdsk_cocos, dd_cocos=dd_cocos, cocos_clockwise_phi=cocos_clockwise_phi,
             add_derived=add_derived,
         )
@@ -322,8 +338,7 @@ end
 """
     geqdsk2imas!(
         g::GEQDSKFile,
-        dd::IMASdd.dd,
-        time_index::Int;
+        dd::IMASdd.dd;
         geqdsk_cocos::Int=0,
         dd_cocos::Int=11,
         cocos_clockwise_phi::Bool=false,
@@ -334,16 +349,14 @@ Utility for writing equilibrium data from a GEQDSKFile to IMAS DD.
 
  * g: GEQDSKFile instance
  * dd: IMAS data dictionary
- * time_index: index of the destination for equilibrium data within equilibrium.time_slice
  * geqdsk_cocos: coordinate convention indentifier of the incoming GEQDSK set. 0 to auto detect.
  * dd_cocos: coordinate convention identifier of the DD. Should normally be left at 11.
  * cocos_clockwise_phi: hint for deciding COCOS, because GEQDSK probably doesn't constrain this.
  * add_derived: switch to use functions within EFIT.jl to extend equilibrium data
 """
 function geqdsk2imas!(
     g::GEQDSKFile,
-    dd::IMASdd.dd,
-    time_index::Int;
+    dd::IMASdd.dd;
     geqdsk_cocos::Int=0,
     dd_cocos::Int=11,
     cocos_clockwise_phi::Bool=false,
@@ -353,7 +366,7 @@ function geqdsk2imas!(
         dd.dataset_description.data_entry.pulse = parse(Int, split(split(g.file, ".")[end-1], "g")[2])
     end
     geqdsk2imas!(
-        g, dd.equilibrium, time_index;
+        g, dd.equilibrium;
         wall=dd.wall, geqdsk_cocos=geqdsk_cocos, dd_cocos=dd_cocos, cocos_clockwise_phi=cocos_clockwise_phi,
         add_derived=add_derived,
     )
@@ -362,8 +375,7 @@ end
 """
     geqdsk2imas!(
         g::GEQDSKFile,
-        eq::IMASdd.equilibrium,
-        time_index::Int;
+        eq::IMASdd.equilibrium;
         wall=nothing,
         geqdsk_cocos::Int=0,
         dd_cocos::Int=11,
@@ -375,7 +387,6 @@ Utility for writing equilibrium data from a GEQDSKFile to IMAS equilibrium IDS
 
  * g: GEQDSKFile instance
  * eq: Reference to equilibrium IDS
- * time_index: index of the destination for equilibrium data within equilibrium.time_slice
  * wall: Optional reference to wall IDS
  * geqdsk_cocos: coordinate convention indentifier of the incoming GEQDSK set. 0 to auto detect.
  * dd_cocos: coordinate convention identifier of the DD. Should normally be left at 11.
@@ -384,8 +395,7 @@ Utility for writing equilibrium data from a GEQDSKFile to IMAS equilibrium IDS
 """
 function geqdsk2imas!(
     g::GEQDSKFile,
-    eq::IMASdd.equilibrium,
-    time_index::Int;
+    eq::IMASdd.equilibrium;
     wall=nothing,
     geqdsk_cocos::Int=0,
     dd_cocos::Int=11,
@@ -400,35 +410,18 @@ function geqdsk2imas!(
     end
     tc = CoordinateConventions.transform_cocos(geqdsk_cocos, dd_cocos)
 
-    # Make sure top level stuff is defined
-    if ismissing(eq, :time)
-        eq.time = Array{Float64}(undef, time_index)
-    end
-    if time_index > length(eq.time)
-        resize!(eq.time, time_index)
-    end
-    if ismissing(eq.vacuum_toroidal_field, :b0)
-        eq.vacuum_toroidal_field.b0 = Array{Float64}(undef, time_index)
-    end
-    if time_index > length(eq.vacuum_toroidal_field.b0)
-        resize!(eq.vacuum_toroidal_field.b0, time_index)
-    end
-
-    # Write top level data
-    eq.time[time_index] = g.time
-    eq.vacuum_toroidal_field.b0[time_index] = g.bcentr .* tc["B"]
-    eq.vacuum_toroidal_field.r0 = g.rcentr
-
-    # Handle time slice data
-    if time_index > length(eq.time_slice)
-        resize!(eq.time_slice, time_index)
+    dd = IMAS.top_dd(eq)
+    original_global_time = dd.global_time
+    try
+        dd.global_time = g.time
+        geqdsk2imas!(
+            g, eq.time_slice[];
+            wall=wall, geqdsk_cocos=geqdsk_cocos, dd_cocos=dd_cocos, cocos_clockwise_phi=cocos_clockwise_phi,
+            add_derived=add_derived,
+        )
+    finally
+        dd.global_time = original_global_time
     end
-    eqt = eq.time_slice[time_index]
-    geqdsk2imas!(
-        g, eqt;
-        wall=wall, geqdsk_cocos=geqdsk_cocos, dd_cocos=dd_cocos, cocos_clockwise_phi=cocos_clockwise_phi,
-        add_derived=add_derived,
-    )
 end
 
 """
@@ -647,27 +640,27 @@ function imas2geqdsk(
     dd_cocos::Int=11,
 )::Vector{GEQDSKFile}
     nt = length(dd.equilibrium.time_slice)
-    return [imas2geqdsk(dd, time_index, geqdsk_cocos=geqdsk_cocos, dd_cocos=dd_cocos) for time_index in 1:nt]
+    return [imas2geqdsk(dd, time, geqdsk_cocos=geqdsk_cocos, dd_cocos=dd_cocos) for time in dd.equilibrium.time]
 end
 
 """
     imas2geqdsk(
         dd::IMASdd.dd,
-        time_index::Int;
+        time::Float64;
         geqdsk_cocos::Int=1,
         dd_cocos::Int=11,
     )::GEQDSKFile
 
 Utility for outputting IMAS equilibrium data from a single time slice to a GEQDSKFile
 
  * dd: IMAS data dictionary instance containing equilibrium data
- * time_index: index of the time slice to output
+ * time: Time of the slice to output in s. Lookup using IMAS global time utilities.
  * geqdsk_cocos: coordinate convention indentifier of the GEQDSKs to be written
  * dd_cocos: coordinate convention identifier of the DD. Should normally be left at 11.
 """
 function imas2geqdsk(
     dd::IMASdd.dd,
-    time_index::Int;
+    time::Float64;
     geqdsk_cocos::Int=1,
     dd_cocos::Int=11,
 )::GEQDSKFile
@@ -679,60 +672,66 @@ function imas2geqdsk(
         shot = 0
     end
 
-    eqt = dd.equilibrium.time_slice[time_index]
-    limiter = dd.wall.description_2d[1].limiter
-    rlim = limiter.unit[1].outline.r ./ tc["R"]
-    zlim = limiter.unit[1].outline.z ./ tc["Z"]
-    limitr = length(rlim)
-
-    # Global and boundary
-    gq = eqt.global_quantities
-    rmaxis = gq.magnetic_axis.r / tc["R"]
-    zmaxis = gq.magnetic_axis.z / tc["Z"]
-    current = gq.ip / tc["I"]
-    simag = gq.psi_axis / tc["PSI"]
-    sibry = gq.psi_boundary / tc["PSI"]
-    rcentr = eqt.boundary.geometric_axis.r / tc["R"]
-    zmid = eqt.boundary.geometric_axis.z / tc["Z"]
-    rbbbs = eqt.boundary.outline.r ./ tc["R"]
-    zbbbs = eqt.boundary.outline.z ./ tc["Z"]
-    nbbbs = length(rbbbs)
-
-    # 1D profiles
-    p1 = eqt.profiles_1d
-    psi = p1.psi ./ tc["PSI"]
-    qpsi = p1.q ./ tc["Q"]
-    pres = p1.pressure ./ tc["P"]
-    pprime = p1.dpressure_dpsi ./ tc["PPRIME"]
-    fpol = p1.f ./ tc["F"]
-    ffprim = p1.f_df_dpsi ./ tc["F_FPRIME"]
-    rhovn = p1.rho_tor_norm
-
-    # 2D flux map
-    p2 = eqt.profiles_2d[1]
-    r = p2.grid.dim1 ./ tc["R"]
-    z = p2.grid.dim2 ./ tc["Z"]
-    psirz = p2.psi ./ tc["PSI"]
-
-    bcentr = dd.equilibrium.vacuum_toroidal_field.b0[time_index] ./ tc["B"]
-    time = eqt.time
-    @debug "eqt.time = $(eqt.time), time out = $time"
-    rleft = minimum(r)
-    rdim = maximum(r) - rleft
-    zdim = maximum(z) - minimum(z)
-    nw = length(r)
-    nh = length(z)
-    itime = Int(round(time * 1000))
-    gfile = "g" * lpad(shot, 6, '0') * "." * lpad(itime, 5, '0')
-    @debug "shot $shot, gfile = $gfile"
-
-    r = range(rleft, rleft + rdim, length=nw)
-    z = range(zmid - 0.5*zdim, zmid + 0.5*zdim, length=nh)
-    psi = range(simag, sibry, length=nw)
-
-    return GEQDSKFile(
-        gfile,time,nw,nh,r,z,rdim,zdim,rleft,zmid,nbbbs,rbbbs,zbbbs,limitr,rlim,zlim,
-        rcentr,bcentr,rmaxis,zmaxis,simag,sibry,psi,current,fpol,pres,ffprim,pprime,
-        qpsi,psirz,rhovn,
-    )
+    original_global_time = dd.global_time
+    try
+        dd.global_time = time
+        eqt = dd.equilibrium.time_slice[]
+        limiter = dd.wall.description_2d[1].limiter
+        rlim = limiter.unit[1].outline.r ./ tc["R"]
+        zlim = limiter.unit[1].outline.z ./ tc["Z"]
+        limitr = length(rlim)
+
+        # Global and boundary
+        gq = eqt.global_quantities
+        rmaxis = gq.magnetic_axis.r / tc["R"]
+        zmaxis = gq.magnetic_axis.z / tc["Z"]
+        current = gq.ip / tc["I"]
+        simag = gq.psi_axis / tc["PSI"]
+        sibry = gq.psi_boundary / tc["PSI"]
+        rcentr = eqt.boundary.geometric_axis.r / tc["R"]
+        zmid = eqt.boundary.geometric_axis.z / tc["Z"]
+        rbbbs = eqt.boundary.outline.r ./ tc["R"]
+        zbbbs = eqt.boundary.outline.z ./ tc["Z"]
+        nbbbs = length(rbbbs)
+
+        # 1D profiles
+        p1 = eqt.profiles_1d
+        psi = p1.psi ./ tc["PSI"]
+        qpsi = p1.q ./ tc["Q"]
+        pres = p1.pressure ./ tc["P"]
+        pprime = p1.dpressure_dpsi ./ tc["PPRIME"]
+        fpol = p1.f ./ tc["F"]
+        ffprim = p1.f_df_dpsi ./ tc["F_FPRIME"]
+        rhovn = p1.rho_tor_norm
+
+        # 2D flux map
+        p2 = eqt.profiles_2d[1]
+        r = p2.grid.dim1 ./ tc["R"]
+        z = p2.grid.dim2 ./ tc["Z"]
+        psirz = p2.psi ./ tc["PSI"]
+
+        bcentr = dd.equilibrium.vacuum_toroidal_field.b0[] ./ tc["B"]
+        time = eqt.time
+        @debug "eqt.time = $(eqt.time), time out = $time"
+        rleft = minimum(r)
+        rdim = maximum(r) - rleft
+        zdim = maximum(z) - minimum(z)
+        nw = length(r)
+        nh = length(z)
+        itime = Int(round(time * 1000))
+        gfile = "g" * lpad(shot, 6, '0') * "." * lpad(itime, 5, '0')
+        @debug "shot $shot, gfile = $gfile"
+
+        r = range(rleft, rleft + rdim, length=nw)
+        z = range(zmid - 0.5*zdim, zmid + 0.5*zdim, length=nh)
+        psi = range(simag, sibry, length=nw)
+
+        return GEQDSKFile(
+            gfile,time,nw,nh,r,z,rdim,zdim,rleft,zmid,nbbbs,rbbbs,zbbbs,limitr,rlim,zlim,
+            rcentr,bcentr,rmaxis,zmaxis,simag,sibry,psi,current,fpol,pres,ffprim,pprime,
+            qpsi,psirz,rhovn,
+        )
+    finally
+        dd.global_time = original_global_time
+    end
 end
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -149,14 +149,15 @@ end
     # Add one geqdsk
     println("test file 1 info: filename = $(gs[1].file), time = $(gs[1].time)")
     dd = IMASdd.dd()
-    eqt = resize!(dd.equilibrium.time_slice, 1)[1]
+    eqt = resize!(dd.equilibrium.time_slice, length(gs))[1]
+    dd.equilibrium.time = [g.time, for g in gs]
     EFIT.geqdsk2imas!(gs[1], eqt; wall=dd.wall, add_derived=true, cocos_clockwise_phi=cocos_clockwise_phi)
     @test length(eqt.profiles_2d[1].grid.dim1) > 1
     @test eqt.time == gs[1].time
     # Add another geqdsk to another index
     println("test file 2 info: filename = $(gs[2].file), time = $(gs[2].time)")
     idx = 2
-    EFIT.geqdsk2imas!(gs[2], dd.equilibrium, idx, add_derived=true, cocos_clockwise_phi=cocos_clockwise_phi)
+    EFIT.geqdsk2imas!(gs[2], dd.equilibrium, add_derived=true, cocos_clockwise_phi=cocos_clockwise_phi)
     eqt2 = dd.equilibrium.time_slice[idx]
     @test length(eqt2.profiles_2d[1].grid.dim1) > 1
     @test abs(dd.equilibrium.vacuum_toroidal_field.b0[idx]) > 0.0