You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
111 lines
3.2 KiB
Python
111 lines
3.2 KiB
Python
# %%
|
|
import pandas as pd
|
|
|
|
from itertools import product
|
|
|
|
|
|
def get_state_vect_cols(prefix=''):
|
|
if prefix:
|
|
prefix += '_'
|
|
vectors = ['r', 'v']
|
|
components = ['x', 'y', 'z']
|
|
col_names = [f'{prefix}{v}_{c}' for v, c in product(vectors, components)]
|
|
return col_names
|
|
|
|
|
|
# %%
|
|
df = pd.read_parquet("traindata/physics_preds.parquet")
|
|
test_set = df[df['aso_id'] == "05277"]
|
|
|
|
train_set = df.groupby('aso_id').apply(lambda x: x.head(x.count()[0] - 3))
|
|
print(df.count()[0], train_set.count()[0], test_set.count()[0])
|
|
|
|
# %%
|
|
from sklearn.model_selection import train_test_split
|
|
|
|
feature_cols = [
|
|
'elapsed_seconds'
|
|
] + get_state_vect_cols('physics_pred') + get_state_vect_cols('start')
|
|
print(feature_cols)
|
|
# The target values are the errors between the physical model predictions
|
|
# and the ground truth observations
|
|
target_cols = get_state_vect_cols('physics_err')
|
|
print(target_cols)
|
|
# Create feature and target matrices
|
|
X = df[feature_cols]
|
|
y = df[target_cols]
|
|
data_keys = ['X_train', 'X_test', 'y_train', 'y_test']
|
|
data_vals = train_test_split(X, y, test_size=0.2)
|
|
train_test_data = dict(zip(data_keys, data_vals))
|
|
|
|
# %%
|
|
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.data import TensorDataset, DataLoader
|
|
from sklearn import metrics
|
|
|
|
TestNet2 = lambda: nn.Sequential(
|
|
nn.Linear(13, 64),
|
|
nn.ReLU6(),
|
|
nn.Linear(64, 64),
|
|
nn.LeakyReLU(),
|
|
nn.Linear(64, 1),
|
|
)
|
|
|
|
nets = {}
|
|
|
|
X_train = torch.tensor(train_test_data["X_train"].values,
|
|
dtype=torch.float32).cuda()
|
|
y_train = train_test_data["y_train"]
|
|
X_test = torch.tensor(train_test_data['X_test'].values,
|
|
dtype=torch.float32).cuda()
|
|
y_test = train_test_data['y_test']
|
|
r2s = []
|
|
for target_col in y_train.columns:
|
|
y1 = torch.tensor(y_train[target_col].values,
|
|
dtype=torch.float32).reshape(-1, 1).cuda()
|
|
print(X.shape, y1.shape)
|
|
net = TestNet2().cuda()
|
|
opti = torch.optim.SGD(net.parameters(), lr=0.02)
|
|
loss_func = nn.MSELoss()
|
|
train_dataloader = DataLoader(TensorDataset(X_train, y1), batch_size=320)
|
|
for t in range(10000):
|
|
for batch, (x, y) in enumerate(train_dataloader):
|
|
pred = net(x)
|
|
loss = loss_func(pred, y)
|
|
opti.zero_grad()
|
|
torch.sqrt(loss).backward()
|
|
opti.step()
|
|
with torch.no_grad():
|
|
y = y_test[target_col]
|
|
y_hat = net(X_test).cpu().numpy()
|
|
rmse = metrics.mean_squared_error(y, y_hat, squared=False)
|
|
r2 = metrics.r2_score(y, y_hat)
|
|
r2s.append(r2)
|
|
print(f"Epoch {t}: r2={r2}, rmse={rmse}")
|
|
nets[target_col] = net
|
|
print(target_col)
|
|
break
|
|
|
|
import matplotlib.pyplot as plt
|
|
|
|
plt.plot(r2s)
|
|
plt.show()
|
|
|
|
# %%
|
|
X, ys = train_test_data['X_test'], train_test_data['y_test']
|
|
evals = []
|
|
with torch.no_grad():
|
|
for target_col, net in nets.items():
|
|
y_hat = net(torch.tensor(X.values)) # fake
|
|
y_hat = y_hat.detach().numpy()
|
|
y = ys[target_col] # real
|
|
print(y)
|
|
print(y_hat)
|
|
rmse = metrics.mean_squared_error(y, y_hat, squared=False)
|
|
r2 = metrics.r2_score(y, y_hat)
|
|
eval_dict = {'Error': target_col, 'RMSE': rmse, 'R^2': r2}
|
|
evals.append(eval_dict)
|
|
print(pd.DataFrame(evals))
|
|
# %%
|