From 325e5bd7e5cf14d8515025687d7eb9080189b108 Mon Sep 17 00:00:00 2001
From: Nuwan Yapa <yapanuwan@gmail.com>
Date: Wed, 10 Jul 2024 17:15:11 -0400
Subject: [PATCH] CSV export to one file

---
 calculations/3body_Berggren_B2R_EC.jl |  4 +---
 calculations/3body_HO_B2R_EC.jl       |  4 +---
 calculations/3body_HO_R2R_EC.jl       |  4 +---
 calculations/B2R_Berggren_poles.jl    | 12 +++---------
 helper.jl                             | 22 ++++++++++++++++++++--
 5 files changed, 26 insertions(+), 20 deletions(-)

diff --git a/calculations/3body_Berggren_B2R_EC.jl b/calculations/3body_Berggren_B2R_EC.jl
index 5ade9cd..dd2a290 100644
--- a/calculations/3body_Berggren_B2R_EC.jl
+++ b/calculations/3body_Berggren_B2R_EC.jl
@@ -128,9 +128,7 @@ for c in extrapolating_c
     push!(extrapolated, nearest(evals, current_E))
 end
 
-open("temp/Berggren_B2R.training.csv", "w") do f; writedlm(f, hcat(reim(training)...)); end
-open("temp/Berggren_B2R.exact.csv", "w") do f; writedlm(f, hcat(reim(exact)...)); end
-open("temp/Berggren_B2R.extrapolated.csv", "w") do f; writedlm(f, hcat(reim(extrapolated)...)); end
+exportCSV("temp/Berggren_B2R.csv", (training, exact, extrapolated), ("training", "exact", "extrapolated"))
 
 scatter(real.(training),imag.(training), label="training")
 scatter!(real.(exact),imag.(exact), label="exact")
diff --git a/calculations/3body_HO_B2R_EC.jl b/calculations/3body_HO_B2R_EC.jl
index f91b50e..5312305 100644
--- a/calculations/3body_HO_B2R_EC.jl
+++ b/calculations/3body_HO_B2R_EC.jl
@@ -85,9 +85,7 @@ for c in extrapolating_c
     push!(extrapolated, nearest(evals, current_E))
 end
 
-open("temp/HO_B2R.training.csv", "w") do f; writedlm(f, hcat(reim(training)...)); end
-open("temp/HO_B2R.exact.csv", "w") do f; writedlm(f, hcat(reim(exact)...)); end
-open("temp/HO_B2R.extrapolated.csv", "w") do f; writedlm(f, hcat(reim(extrapolated)...)); end
+exportCSV("temp/HO_B2R.csv", (training, exact, extrapolated), ("training", "exact", "extrapolated"))
 
 scatter(real.(training),imag.(training), label="training")
 scatter!(real.(exact),imag.(exact), label="exact")
diff --git a/calculations/3body_HO_R2R_EC.jl b/calculations/3body_HO_R2R_EC.jl
index 1396c56..d71e3bd 100644
--- a/calculations/3body_HO_R2R_EC.jl
+++ b/calculations/3body_HO_R2R_EC.jl
@@ -46,9 +46,7 @@ for c in extrapolating_c
     push!(extrapolated, nearest(evals, current_E))
 end
 
-open("temp/NCSM.training.csv", "w") do f; writedlm(f, hcat(reim(training)...)); end
-open("temp/NCSM.exact.csv", "w") do f; writedlm(f, hcat(reim(exact)...)); end
-open("temp/NCSM.extrapolated.csv", "w") do f; writedlm(f, hcat(reim(extrapolated)...)); end
+exportCSV("temp/NCSM.csv", (training, exact, extrapolated), ("training", "exact", "extrapolated"))
 
 scatter(real.(training),imag.(training), label="training")
 scatter!(real.(exact),imag.(exact), label="exact")
diff --git a/calculations/B2R_Berggren_poles.jl b/calculations/B2R_Berggren_poles.jl
index f3c3e86..0d97164 100644
--- a/calculations/B2R_Berggren_poles.jl
+++ b/calculations/B2R_Berggren_poles.jl
@@ -1,4 +1,5 @@
-using Plots, DelimitedFiles
+using Plots
+include("../helper.jl")
 include("../p_space.jl")
 
 # contour
@@ -51,14 +52,7 @@ for (j, c) in enumerate(extrapolate_points)
     extrapolate_E[j] = evals[i]
 end
 
-open("temp/2b_GSM.csv", "w") do f
-    write(f, "# training \n")
-    writedlm(f, hcat(reim(training_E)...))
-    write(f, "# exact \n")
-    writedlm(f, hcat(reim(exact_E)...))
-    write(f, "# extrapolate \n")
-    writedlm(f, hcat(reim(extrapolate_E)...))
-end
+exportCSV("temp/2b_GSM.csv", (training_E, exact_E, extrapolate_E), ("training", "exact", "extrapolated"))
 
 scatter(real.(training_E), imag.(training_E), label="training")
 scatter!(real.(exact_E), imag.(exact_E), label="exact")
diff --git a/helper.jl b/helper.jl
index dbd70a5..4d284ba 100644
--- a/helper.jl
+++ b/helper.jl
@@ -1,4 +1,4 @@
-using LinearAlgebra
+using LinearAlgebra, DelimitedFiles
 
 "Sum over array while minimizing catastrophic cancellation as much as possible"
 function better_sum(arr::Array{Float64})
@@ -29,4 +29,22 @@ function kron_sum(A::AbstractMatrix, B::AbstractMatrix)
 end
 
 "Flattened vector version of Iterators.product(...) with index hierachy reversed -- leftmost index has the highest hierachy"
-iter_prod(args...) = reverse.(collect(Iterators.product(reverse(args)...))[:])
\ No newline at end of file
+iter_prod(args...) = reverse.(collect(Iterators.product(reverse(args)...))[:])
+
+"Export CSV data for a scatter plot taking in data as a list of complex vectors (x=Re, y=Im), and a list of corresponding labels, typically used for EC results"
+function exportCSV(file::String, data, labels=nothing)
+    columns = ["x" "y" "label"]
+
+    if isnothing(labels)
+        columns[end] = ""
+        labels = fill("", length(data))
+    end
+    
+    open(file, "w") do f
+        writedlm(f, columns)
+        for (d, label) in zip(data, labels)
+            l = fill(label, length(d))
+            writedlm(f, hcat(reim(d)..., l))
+        end
+    end
+end
\ No newline at end of file