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

Merge branch 'main' into pmm/k

This commit is contained in:
Nuwan Yapa 2025-04-27 20:43:21 -04:00
parent ae5b4ecea8 3d7b964cde
commit b2c2b80a83
1 changed files with 13 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
calculations/PMM.py
View File

@ -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()