diff --git a/calculations/PMM.py b/calculations/PMM.py index b3c1875..85fa238 100644 --- a/calculations/PMM.py +++ b/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()