yapanuwan
/
BergEC-jl
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Emulate k instead of E

This commit is contained in:
Nuwan Yapa 2025-04-27 17:42:37 -04:00
parent 65d0f44b69
commit ae5b4ecea8
1 changed files with 37 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

53
calculations/PMM.py
View File

@ -1,15 +1,19 @@
#%% #%%
import pandas as pd import pandas as pd
import torch import torch
import numpy as np
#%% #%%
df = pd.read_csv('../temp/2body_data.csv') 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['E'] = df['re_E'] + 1j * df['im_E']
df['k'] = np.sqrt(df['E'])
train_data = df[df['re_E'] < 0] train_data = df[df['re_E'] < 0]
target_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 = 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 # hyperparameters
@ -27,14 +31,14 @@ H1 = (H1 + torch.transpose(H1, 0, 1)).requires_grad_() # symmetric
lr = 0.05 lr = 0.05
epochs = 100000 epochs = 100000
for epoch in range(epochs): for epoch in range(epochs):
Es = torch.empty(len(train_data), dtype=torch.complex128) ks = torch.empty(len(train_data), dtype=torch.complex128)
for (index, (c, E)) in enumerate(zip(train_cs, train_Es)): for (index, (c, k)) in enumerate(zip(train_cs, train_ks)):
H = H0 + c * H1 H = H0 + c * H1
evals = torch.linalg.eigvals(H) evals = torch.linalg.eigvals(H)
i = torch.argmin(torch.abs(evals - E)) # TODO: more robust way to identify the eigenvector i = torch.argmin(torch.abs(evals - k)) # TODO: more robust way to identify the eigenvector
Es[index]= evals[i] ks[index]= evals[i]
loss = ((Es - train_Es).abs() ** 2).sum() loss = ((ks - train_ks).abs() ** 2).sum()
if epoch % 1000 == 0: if epoch % 1000 == 0:
print(f"Training {(epoch+1)/epochs:.1%} \t Loss: {loss}") print(f"Training {(epoch+1)/epochs:.1%} \t Loss: {loss}")
@ -52,21 +56,38 @@ for epoch in range(epochs):
# %% # %%
# evaluate for all points # evaluate for all points
all_c = torch.tensor(df['c'].values, dtype=torch.float64) all_c = torch.tensor(df['c'].values, dtype=torch.float64)
exact_E = torch.tensor(df['E'].values, dtype=torch.complex128) exact_k = torch.tensor(df['k'].values, dtype=torch.complex128)
pred_Es = torch.empty(len(df), dtype=torch.complex128) pred_ks = np.empty(len(df), dtype=np.complex128)
with torch.no_grad(): 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 H = H0 + c * H1
evals = torch.linalg.eigvals(H) evals = torch.linalg.eigvals(H)
i = torch.argmin(torch.abs(evals - E)) # TODO: more robust way to identify the eigenvector i = torch.argmin(torch.abs(evals - k)) # TODO: more robust way to identify the eigenvector
pred_Es[index]= evals[i] pred_ks[index]= evals[i]
pred_Es = pred_ks ** 2
# %% # %%
# plot the results # plot the results
import matplotlib.pyplot as plt 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()