From 65d0f44b698a9fb6576163de6bcd7c660480ea8b Mon Sep 17 00:00:00 2001
From: Nuwan Yapa <yapanuwan@gmail.com>
Date: Sat, 26 Apr 2025 21:11:49 -0400
Subject: [PATCH] Seperate Julia file for exporting PMM data

---
 calculations/2body_data.jl | 33 +++++++++++++++++++++++++++++++++
 calculations/PMM.jl        |  6 ------
 calculations/PMM.py        | 10 +++++-----
 3 files changed, 38 insertions(+), 11 deletions(-)
 create mode 100644 calculations/2body_data.jl

diff --git a/calculations/2body_data.jl b/calculations/2body_data.jl
new file mode 100644
index 0000000..763c8db
--- /dev/null
+++ b/calculations/2body_data.jl
@@ -0,0 +1,33 @@
+using Roots, DelimitedFiles
+
+include("../EC.jl")
+include("../common.jl")
+include("../p_space.jl")
+
+μ = 0.5
+V_system(c) = (p, q) -> c*(-5*g0(sqrt(3), p, q) + 2*g0(sqrt(10), p, q)) # ResonanceEC: Eq. (20)
+
+# determining c0 with EC
+temp_c = range(1.1, 0.9, 3)
+p, w = get_mesh([0, 8], [256])
+H0 = get_T_matrix(p, μ)
+V = get_V_matrix(V_system(1), p, w)
+EC = affine_EC(H0, V, w)
+train!(EC, temp_c; ref_eval=-0.2, CAEC=false, verbose=false)
+quick_extrapolate(c) = minimum(abs2, get_extrapolated_evals(EC.H0_EC, EC.H1_EC, EC.N_EC, c, 0))
+c0 = find_zero(quick_extrapolate, 0.85)
+
+training_c = range(1.2, 0.9, 9) # original: range(1.35, 0.9, 5)
+extrapolating_c = range(0.78, 0.45, 7) # original: range(0.75, 0.40, 8)
+
+data_c = vcat(training_c, extrapolating_c)
+data_E = [quick_pole_E(V_system(c)) for c in data_c]
+
+# export to CSV
+file = "temp/2body_data.csv"
+delim = ','
+open(file, "w") do f
+    writedlm(f, ["c" "re_E" "im_E"], delim)
+    writedlm(f, [c0 0 0], delim) # first entry for the threshold
+    writedlm(f, hcat(data_c, real.(data_E), imag.(data_E)), delim)
+end
diff --git a/calculations/PMM.jl b/calculations/PMM.jl
index 3f86d54..f408738 100644
--- a/calculations/PMM.jl
+++ b/calculations/PMM.jl
@@ -69,9 +69,3 @@ end
 scatter(real.(data_E), imag.(data_E), label="training", title="PMM", xlabel="Re E", ylabel="Im E")
 scatter!(real.(exact_E), imag.(exact_E), label="exact", m=:+)
 scatter!(real.(extrapolated_E), imag.(extrapolated_E), label="predicted", m=:x)
-
-# export results
-using DataFrames, CSV
-df = DataFrame(c=all_c, re_E=real.(exact_E), im_E=imag.(exact_E))
-#pushfirst!(df, (c0, 0, 0)) # add the threshold
-CSV.write("temp/PMM.csv", df)
diff --git a/calculations/PMM.py b/calculations/PMM.py
index 69cac24..b3c1875 100644
--- a/calculations/PMM.py
+++ b/calculations/PMM.py
@@ -3,13 +3,13 @@ import pandas as pd
 import torch
 
 #%%
-df = pd.read_csv('../temp/PMM.csv')
+df = pd.read_csv('../temp/2body_data.csv')
 df['E'] = df['re_E'] + 1j * df['im_E']
-train_data = df[abs(df['im_E']) < 1e-5]
-target_data = df[abs(df['im_E']) > 1e-5]
+train_data = df[df['re_E'] < 0]
+target_data = df[df['re_E'] > 0]
 
-train_cs = torch.tensor(train_data['c'], dtype=torch.float64)
-train_Es = torch.tensor(train_data['E'], dtype=torch.complex128)
+train_cs = torch.tensor(train_data['c'].to_numpy(), dtype=torch.float64)
+train_Es = torch.tensor(train_data['E'].to_numpy(), dtype=torch.complex128)
 
 #%%
 # hyperparameters