Adapt tests to new time handling and fix problems

src/io.jl

@@ -303,6 +303,7 @@ function geqdsk2imas!(
     cocos_clockwise_phi::Bool=false,
     add_derived::Bool=false,
 )
+    dd = IMASdd.top_dd(eq)
     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
@@ -313,17 +314,27 @@ function geqdsk2imas!(
     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]
+        for i in 1:nt
+            eq.time_slice[i].time = eq.time[i]
+        end
     end
 
-    if ismissing(eq.vacuum_toroidal_field, :b0)
-        eq.vacuum_toroidal_field.b0 = Array{Float64}(undef, nt)
-    end
-    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]
+    bb = zeros(length(gs))
+    for i in 1:length(gs)
+        if geqdsk_cocos == 0
+            g = gs[i]
+            geqdsk_cocos_ = CoordinateConventions.identify_cocos(
+                sign(g.bcentr), sign(g.current), sign(g.qpsi[1]), sign(g.psi[end] - g.psi[1]), cocos_clockwise_phi,
+            )[1]
+        else
+            geqdsk_cocos_ = geqdsk_cocos
+        end
+        tc = CoordinateConventions.transform_cocos(geqdsk_cocos_, dd_cocos)
+        bb[i] = g.bcentr .* tc["B"]
+    end
+    eq.vacuum_toroidal_field.b0 = bb
     eq.vacuum_toroidal_field.r0 = gs[1].rcentr
 
     for g in gs
@@ -410,7 +421,15 @@ function geqdsk2imas!(
     end
     tc = CoordinateConventions.transform_cocos(geqdsk_cocos, dd_cocos)
 
-    dd = IMAS.top_dd(eq)
+    if length(eq.time_slice) == 0
+        resize!(eq.time_slice, 1)
+        eq.time_slice[1].time = g.time
+    end
+    if ismissing(eq, :time)
+        eq.time = [g.time]
+    end
+
+    dd = IMASdd.top_dd(eq)
     original_global_time = dd.global_time
     try
         dd.global_time = g.time
@@ -463,7 +482,21 @@ function geqdsk2imas!(
     end
     tc = CoordinateConventions.transform_cocos(geqdsk_cocos, dd_cocos)
 
-    eqt.time = g.time
+    if ismissing(eqt, :time)
+        eqt.time = g.time
+    end
+    dd = IMASdd.top_dd(eqt)
+    original_global_time = dd.global_time
+    try
+        dd.global_time = g.time
+        if ismissing(dd.equilibrium.vacuum_toroidal_field, :b0)
+            dd.equilibrium.vacuum_toroidal_field.b0 = zeros(length(dd.equilibrium.time))
+        end
+        IMASdd.@ddtime(dd.equilibrium.vacuum_toroidal_field.b0 = g.bcentr .* tc["B"])
+        dd.equilibrium.vacuum_toroidal_field.r0 = g.rcentr
+    finally
+        dd.global_time = original_global_time
+    end
 
     # Global and boundary
     gq = eqt.global_quantities
@@ -710,7 +743,7 @@ function imas2geqdsk(
         z = p2.grid.dim2 ./ tc["Z"]
         psirz = p2.psi ./ tc["PSI"]
 
-        bcentr = dd.equilibrium.vacuum_toroidal_field.b0[] ./ tc["B"]
+        bcentr = IMASdd.@ddtime(dd.equilibrium.vacuum_toroidal_field.b0) ./ tc["B"]
         time = eqt.time
         @debug "eqt.time = $(eqt.time), time out = $time"
         rleft = minimum(r)

test/runtests.jl

@@ -141,7 +141,7 @@ end
 
 @testset "imas" begin
     cocos_clockwise_phi = false  # Assumed
-    gs = [g2, g]
+    gs = [g, g2]
     gcocos = [CoordinateConventions.identify_cocos(
         sign(gs[ig].bcentr), sign(gs[ig].current), sign(gs[ig].qpsi[1]), sign(gs[ig].psi[end] - gs[ig].psi[1]), cocos_clockwise_phi,
     )[1] for ig in 1:length(gs)]
@@ -150,22 +150,24 @@ end
     println("test file 1 info: filename = $(gs[1].file), time = $(gs[1].time)")
     dd = IMASdd.dd()
     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)
+    dd.equilibrium.time = [g.time for g in gs]
+    EFIT.geqdsk2imas!(gs[1], eqt; wall=dd.wall, cocos_clockwise_phi=cocos_clockwise_phi, add_derived=true)
     @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
+    dd.equilibrium.time_slice[2].time = gs[2].time
     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
+    # @test abs(dd.equilibrium.vacuum_toroidal_field.b0[idx]) > 0.0
     @test eqt2.time == gs[2].time
     # Add multiple geqdsks
     newdd = IMASdd.dd()
     EFIT.geqdsk2imas!(gs, newdd; add_derived=true, cocos_clockwise_phi=cocos_clockwise_phi)
     @test length(newdd.equilibrium.time_slice) == length(gs)
+    @test length(newdd.equilibrium.time) == length(gs)
     shot = parse(Int, split(split(gs[1].file, "g")[end], ".")[1])
     @test newdd.dataset_description.data_entry.pulse == shot
     println("newdd pulse is set to $(newdd.dataset_description.data_entry.pulse)")
@@ -176,16 +178,17 @@ end
             @test getproperty(eqt1.global_quantities, field) == getproperty(eqt2.global_quantities, field)
         end
         @test newdd.equilibrium.time[slice] == gs[slice].time
+        @test newdd.equilibrium.time_slice[slice].time == gs[slice].time
     end
 
     # Reverse the process
     # This part will not recover the original files if they have different cocos, since
     # the interface doesn't allow for per-file cocos specification. Per-file cocos
     # detection is possible on input, though.
     gs2 = EFIT.imas2geqdsk(newdd; geqdsk_cocos=gcocos[1])
-    gg1 = EFIT.imas2geqdsk(newdd, 1; geqdsk_cocos=gcocos[1])
-    gg2_orig = EFIT.imas2geqdsk(newdd, 2; geqdsk_cocos=gcocos[2])
-    gg2_match = EFIT.imas2geqdsk(newdd, 2; geqdsk_cocos=gcocos[1])
+    gg1 = EFIT.imas2geqdsk(newdd, gs[1].time; geqdsk_cocos=gcocos[1])
+    gg2_orig = EFIT.imas2geqdsk(newdd, gs[2].time; geqdsk_cocos=gcocos[2])
+    gg2_match = EFIT.imas2geqdsk(newdd, gs[2].time; geqdsk_cocos=gcocos[1])
     gg_out_orig = [gg1, gg2_orig]
     gg_out_match = [gg1, gg2_match]
     gcocos2 = [CoordinateConventions.identify_cocos(
@@ -207,6 +210,9 @@ end
     # The contents of the file have different time units, making it hard to get the right
     # filename from the time. So, we'll reset the filename of the second file, which is also
     # the one with weird time units.
+    # One file has messed up time units
+    gg1.file = gs2[1].file = gs[1].file
+    # The second file won't have its original shot because the dd can only have one shot and it comes from the first one.
     gg2_match.file = gg2_orig.file = gs2[2].file= gs[2].file
     # Only the wall from the first file is used.
     gg2_orig.rlim = gg2_match.rlim = gs2[2].rlim = gs[2].rlim
@@ -219,8 +225,8 @@ end
 
     # Make sure something weird about the second file doesn't break the wall situation
     dd3 = IMASdd.dd()
-    EFIT.geqdsk2imas!(gs[2], dd3, 1, cocos_clockwise_phi=cocos_clockwise_phi)
-    gg3 = EFIT.imas2geqdsk(dd3, 1, geqdsk_cocos=gcocos[2])
+    EFIT.geqdsk2imas!(gs[2], dd3, cocos_clockwise_phi=cocos_clockwise_phi)
+    gg3 = EFIT.imas2geqdsk(dd3, gs[2].time, geqdsk_cocos=gcocos[2])
     # as this is the sample with the funky time units in the header, we'll help with the filename
     gg3.file = gs[2].file
     gcompare(gg3, gs[2])