diff --git a/calculations/PMM.py b/calculations/PMM.py
index b3c1875..81b56ef 100644
--- a/calculations/PMM.py
+++ b/calculations/PMM.py
@@ -1,15 +1,19 @@
 #%%
 import pandas as pd
 import torch
+import numpy as np
 
 #%%
 df = pd.read_csv('../temp/2body_data.csv')
+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'])
+
 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_Es = torch.tensor(train_data['E'].to_numpy(), dtype=torch.complex128)
+train_ks = torch.tensor(train_data['k'].to_numpy(), dtype=torch.complex128)
 
 #%%
 # hyperparameters
@@ -27,14 +31,14 @@ H1 = (H1 + torch.transpose(H1, 0, 1)).requires_grad_() # symmetric
 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)):
+    ks = torch.empty(len(train_data), dtype=torch.complex128)
+    for (index, (c, k)) in enumerate(zip(train_cs, train_ks)):
         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] 
+        i = torch.argmin(torch.abs(evals - k)) # TODO: more robust way to identify the eigenvector
+        ks[index]= evals[i]
 
-    loss = ((Es - train_Es).abs() ** 2).sum()
+    loss = ((ks - train_ks).abs() ** 2).sum()
 
     if epoch % 1000 == 0:
         print(f"Training {(epoch+1)/epochs:.1%} \t Loss: {loss}")
@@ -52,21 +56,38 @@ for epoch in range(epochs):
 # %%
 # evaluate for all points
 all_c = torch.tensor(df['c'].values, dtype=torch.float64)
-exact_E = torch.tensor(df['E'].values, dtype=torch.complex128)
-pred_Es = torch.empty(len(df), dtype=torch.complex128)
+exact_k = torch.tensor(df['k'].values, dtype=torch.complex128)
+pred_ks = np.empty(len(df), dtype=np.complex128)
 with torch.no_grad():
-    for (index, (c, E)) in enumerate(zip(all_c, exact_E)):
+    for (index, (c, k)) in enumerate(zip(all_c, exact_k)):
         H = H0 + c * H1
         evals = torch.linalg.eigvals(H)
-        i = torch.argmin(torch.abs(evals - E)) # TODO: more robust way to identify the eigenvector
-        pred_Es[index]= evals[i]
+        i = torch.argmin(torch.abs(evals - k)) # TODO: more robust way to identify the eigenvector
+        pred_ks[index]= evals[i]
+
+pred_Es = pred_ks ** 2
 
 # %%
 # plot the results
 import matplotlib.pyplot as plt
-plt.scatter(train_data['re_E'], train_data['im_E'], label='training')
-plt.scatter(target_data['re_E'], target_data['im_E'], label='target')
-plt.scatter(pred_Es.real, pred_Es.imag, marker='x', label='predicted')
-plt.legend()
 
-# %%
+fig, axs = plt.subplots(2, 1, figsize=(8, 12))  # Create a figure with two vertical panels
+
+# First panel: k values
+axs[0].scatter(np.real(train_data['k']), np.imag(train_data['k']), label='training')
+axs[0].scatter(np.real(target_data['k']), np.imag(target_data['k']), label='target')
+axs[0].scatter(np.real(pred_ks), np.imag(pred_ks), marker='x', label='predicted')
+axs[0].set_xlabel('Re(k)')
+axs[0].set_ylabel('Im(k)')
+axs[0].legend()
+
+# Second panel: E values
+axs[1].scatter(np.real(train_data['E']), np.imag(train_data['E']), label='training')
+axs[1].scatter(np.real(target_data['E']), np.imag(target_data['E']), label='target')
+axs[1].scatter(np.real(pred_Es), np.imag(pred_Es), marker='x', label='predicted')
+axs[1].set_xlabel('Re(E)')
+axs[1].set_ylabel('Im(E)')
+axs[1].legend()
+
+plt.tight_layout()  # Adjust spacing between panels
+plt.show()