Robust way to identify eigenvalue using finer c steps

calculations/PMM.py

@@ -1,14 +1,15 @@
 #%%
 import pandas as pd
 import torch
+import numpy as np
 
 #%%
-df = pd.read_csv('../temp/2body_data.csv')
+df = pd.read_csv('../temp/2body_data.csv').sort_values(by='c')
 df['E'] = df['re_E'] + 1j * df['im_E']
 train_data = df[df['re_E'] < 0]
 target_data = df[df['re_E'] > 0]
 
-train_cs = torch.tensor(train_data['c'].to_numpy(), dtype=torch.float64)
+train_cs = train_data['c'].to_numpy()
 train_Es = torch.tensor(train_data['E'].to_numpy(), dtype=torch.complex128)
 
 #%%
@@ -24,15 +25,23 @@ H1 = (H1 + torch.transpose(H1, 0, 1)).requires_grad_() # symmetric
 #%%
 # training
 
+# generate a set of c values to follow by subdividing the training cs
+subdivisions = 3
+c_steps = np.concatenate([np.linspace(start, stop, subdivisions, endpoint=False) for (start, stop) in zip(train_cs, train_cs[1:])])
+c_steps = np.append(c_steps, train_cs[-1])
+
 lr = 0.05
 epochs = 100000
 for epoch in range(epochs):
     Es = torch.empty(len(train_data), dtype=torch.complex128)
-    for (index, (c, E)) in enumerate(zip(train_cs, train_Es)):
+    current_E = 0.0 # start at the threshold
+    for c in c_steps:
         H = H0 + c * H1
         evals = torch.linalg.eigvals(H)
-        i = torch.argmin(torch.abs(evals - E)) # TODO: more robust way to identify the eigenvector
-        Es[index]= evals[i] 
+        current_E = evals[torch.argmin(torch.abs(evals - current_E))]
+        if np.any(c == train_cs):
+            index = np.where(c == train_cs)[0][0]
+            Es[index] = current_E
 
     loss = ((Es - train_Es).abs() ** 2).sum()