Merge pull request #17 from mgyoo86/feature/write_geqdsk_file

writeg G-EQDSK file

Project.toml

@@ -7,9 +7,11 @@ version = "1.0.2"
 IMASdd = "c5a45a97-b3f9-491c-b9a7-aa88c3bc0067"
 Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
 PolygonOps = "647866c9-e3ac-4575-94e7-e3d426903924"
+Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 
 [compat]
 IMASdd = "2"
 Interpolations = "0.14.7, 0.15"
 PolygonOps = "0.1"
+Printf = "1.11.0"
 julia = "^1.2"

src/EFIT.jl

@@ -3,7 +3,7 @@ __precompile__()
 module EFIT
 
 include("io.jl")
-export readg, GEQDSKFile
+export readg, writeg, GEQDSKFile
 
 include("utils.jl")
 export triangularity, ellipticity, elongation, minor_radius, major_radius, aspect_ratio, elevation

src/io.jl

@@ -1,3 +1,5 @@
+using Printf
+
 mutable struct GEQDSKFile
     file::String                    # Source file
     time::Float64                   # Time of the equilibrium reconstruction, in seconds
@@ -32,6 +34,18 @@ mutable struct GEQDSKFile
     rhovn::Vector{Float64}          # Square root of toroidal flux, normalized
 end
 
+# Define hash function using fieldnames and getfield
+function Base.hash(x::GEQDSKFile, h::UInt)
+    for field in fieldnames(GEQDSKFile)
+        h = hash(getfield(x, field), h)
+    end
+    return h
+end
+# Define == using fieldnames and getfield
+Base.:(==)(a::GEQDSKFile, b::GEQDSKFile) = all(field -> getfield(a, field) == getfield(b, field), fieldnames(GEQDSKFile))
+# Define isequal using fieldnames and getfield
+Base.isequal(a::GEQDSKFile, b::GEQDSKFile) = all(field -> isequal(getfield(a, field), getfield(b, field)), fieldnames(GEQDSKFile))
+
 function Base.show(io::IO, g::GEQDSKFile)
     print(io,"GEQDSKFile: \"",g.file,"\"")
 end
@@ -229,3 +243,65 @@ function readg(gfile; set_time=nothing)
 
     return g
 end
+
+function write_vector_data_in_chunks(io::IOStream, v::Vector{Float64})
+    # Iterate over the vector in chunks of 5 elements
+    for chunk in Iterators.partition(v, 5)
+        # Format each number with %15.8E
+        formatted_numbers = [@sprintf("%15.8E", num) for num in chunk]
+        # Join the formatted numbers into a single string with spaces
+        line = join(formatted_numbers, " ")
+        # Write the line to the io stream
+        println(io, line)
+    end
+end
+
+
+# Function to write the GEQDSKFile struct to a G-file with explicit error handling
+function writeg(g::GEQDSKFile, filename::String; desc::String="description")
+    if isdir(filename)
+        @error("Error: A directory with the name '$filename' already exists.")
+        return false
+    elseif isfile(filename)
+        @warn("Warning: A file with the name '$filename' already exists. It will be overwritten.")
+    end
+
+    try
+        # Open the target file for writing
+        open(filename, "w") do f
+            # Write header
+            # Note: GEQDSKFile.time is converted to ms (to follow readg's notation)
+            # Remove all newline characters in description
+            clean_desc = replace(desc, "\n" => "")
+            @printf(f,"%s  %.3f  %d  %d  %d\n",clean_desc, 1e3*g.time, 0, g.nw, g.nh)
+
+            @printf(f,"%15.8E %15.8E %15.8E %15.8E %15.8E\n", g.rdim, g.zdim, g.rcentr, g.rleft, g.zmid)
+            @printf(f,"%15.8E %15.8E %15.8E %15.8E %15.8E\n", g.rmaxis, g.zmaxis, g.simag, g.sibry, g.bcentr)
+            @printf(f,"%15.8E %15.8E %15.8E %15.8E %15.8E\n", g.current, g.simag, 0.0, g.rmaxis, 0.0)
+            @printf(f,"%15.8E %15.8E %15.8E %15.8E %15.8E\n", g.zmaxis, 0.0, g.sibry, 0.0, 0.0)
+
+            write_vector_data_in_chunks(f, g.fpol)
+            write_vector_data_in_chunks(f, g.pres)
+            write_vector_data_in_chunks(f, g.ffprim)
+            write_vector_data_in_chunks(f, g.pprime)
+
+            write_vector_data_in_chunks(f, vec(g.psirz))
+
+            write_vector_data_in_chunks(f, g.qpsi)
+
+            @printf(f,"%d %d\n", g.nbbbs, g.limitr)
+
+
+            bbbs_rz = Vector(vec(hcat(g.rbbbs, g.zbbbs)'))
+            write_vector_data_in_chunks(f, vec(bbbs_rz))
+
+            lim_rz = Vector(vec(hcat(g.rlim, g.zlim)'))
+            write_vector_data_in_chunks(f, vec(lim_rz))
+        end
+        @info "Successfully wrote GEQDSKFile to '$filename'."
+        return true
+    catch e
+        @error "An error occurred while writing to file '$filename': $e"
+        return false
+    end
+end

test/runtests.jl

@@ -94,3 +94,31 @@ end
         end
     end
 end
+
+@testset "writeg" begin
+
+    for test_gfile in [EFIT.test_gfile, EFIT.test_gfile2]
+
+        ori_g = readg(EFIT.test_gfile)
+
+        # Create and write a tmporary EQDSK file
+        (tmp_filename, ~) = Base.mktemp()
+        writeg(ori_g, tmp_filename; desc="\nHello, can you read me?\n\n")
+
+        # Read the temporary EQDSK file
+        new_g = readg(tmp_filename)
+
+        # Match the "file" name for a fair comparison
+        ori_g.file = "we got the same name for a fair comparison"
+        new_g.file = "we got the same name for a fair comparison"
+
+        # Check if new_g is equivalent to ori_g
+        @test new_g == ori_g
+        @test isequal(new_g, ori_g)
+
+        # Delete temporary file
+        rm(tmp_filename, force=true)
+    end
+
+
+end