diff --git a/calculations/PMM.py b/calculations/PMM.py
index 81b56ef..4123a7e 100644
--- a/calculations/PMM.py
+++ b/calculations/PMM.py
@@ -4,7 +4,7 @@ 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.loc[df['re_E'] < 0, 'im_E'] = 0 # set im_E = 0 for bound states (to avoid square root issues)
 df['E'] = df['re_E'] + 1j * df['im_E']
 df['k'] = np.sqrt(df['E'])
@@ -12,7 +12,7 @@ df['k'] = np.sqrt(df['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_ks = torch.tensor(train_data['k'].to_numpy(), dtype=torch.complex128)
 
 #%%
@@ -28,15 +28,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):
     ks = torch.empty(len(train_data), dtype=torch.complex128)
-    for (index, (c, k)) in enumerate(zip(train_cs, train_ks)):
+    current_k = 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 - k)) # TODO: more robust way to identify the eigenvector
-        ks[index]= evals[i]
+        current_k = evals[torch.argmin(torch.abs(evals - current_k))]
+        if np.any(c == train_cs):
+            index = np.where(c == train_cs)[0][0]
+            ks[index] = current_k
 
     loss = ((ks - train_ks).abs() ** 2).sum()