USER
def load_trajectory(filename, task):
traj = np.load(filename,allow_pickle=True)
if task == 'task_1' or task == 'task_2':
trajectory = traj['trajectory']
energy = traj['energy']
return trajectory, energy
if task == 'task_3':
trajectory = traj['trajectory']
framework = traj['framework']
energy = traj['energy']
return trajectory, framework, energy
# Below, you can find an example of how to load a trajectory. ```trajectory``` contains the particle position, velocity and charge data. ```energy``` contains the energy of the system at various timesteps. The shape of the arrays is as follows:
#
# ```trajectory```: ```(time, n_bodies, [x, y, v_x, v_y, q])```
#
# ```energy```: ```(time, 1)```
# In[4]:
trajectory, energy = load_trajectory('data/task1_2/train/trajectory_0.npz', 'task_1')
print(f'Trajectory shape: {trajectory.shape}, Energy shape: {energy.shape}')
# The following code provides an example of how to visualize a trajectory. Feel free to modify this code, or write your own function. When evaluating your model in task 2 and 3, keep in mind that you are only allowed to use the data at t=0 (the black dots in the visualization).
# In[5]:
x = trajectory[...,0]
y = trajectory[...,1]
plt.figure(figsize=(4,4))
plt.vlines([0,20],0,20)
plt.hlines([0,20],0,20)
plt.scatter(x[0], y[0], c='black')
for i in range(x.shape[1]):
plt.scatter(x[:,i], y[:,i], s=5)
plt.xlim(-1,21)
plt.ylim(-1,21)
plt.show()
plt.figure(figsize=(4,1))
plt.plot(energy)
plt.xlabel('step')
plt.title('System energy over time')
plt.show();
# In[ ]:
# # Task 1
# Implement task 1 below. Feel free to add extra code cells for different components of your implementation.
# In[14]:
# Define distance metrics
def euclidean_distance(x, y):
return torch.sqrt(torch.sum((x - y)**2))
def inverse_distance(x, y):
return 1 / euclidean_distance(x, y)
def create_adjacency_matrix(trajectory, distance_metric, time_step=0):
positions = trajectory[time_step, :, :2]
n_bodies = positions.shape[0]
adjacency_matrix = np.zeros((n_bodies, n_bodies))
for i in range(n_bodies):
for j in range(n_bodies):
if i != j:
position_i = positions[i]%20
position_j = positions[j]%20
adjacency_matrix[i, j] = distance_metric(position_i, position_j)
return torch.tensor(adjacency_matrix, dtype=torch.float32)
# Validate input
def validate_input(X, adjacency_matrix):
# X should be a 2D tensor
assert X.dim() == 2, f"X must be 2D, but got shape {X.shape}"
# The number of nodes should be the same in X and the adjacency matrix
assert X.shape[0] == adjacency_matrix.shape[0] == adjacency_matrix.shape[1], \
f"Mismatch in number of nodes: got {X.shape[0]} nodes in X, but {adjacency_matrix.shape[0]} nodes in adjacency matrix"
# The adjacency matrix should be square
assert adjacency_matrix.shape[0] == adjacency_matrix.shape[1], \
f"Adjacency matrix must be square, but got shape {adjacency_matrix.shape}"
print("All checks passed.")
# In[15]:
trajectory = torch.tensor(trajectory, dtype=torch.float32)
adjacency_matrix = create_adjacency_matrix(trajectory, inverse_distance, time_step=0)
# In[16]:
adjacency_matrix
from torch_geometric.data import Dataset, Data, DataLoader
class MyDataset(Dataset):
def __init__(self, root, filenames, transform=None, pre_transform=None):
self.filenames = filenames
super(MyDataset, self).__init__(root, transform, pre_transform)
@property
def raw_file_names(self):
return self.filenames
def len(self):
return len(self.filenames)
def get(self, idx):
X, y = load_trajectory(self.filenames[idx], 'task_1')
X = torch.tensor(X, dtype=torch.float32)
y = torch.tensor(y, dtype=torch.float32)
data_list = []
for time_step in range(X.shape[0]):
adjacency_matrix = create_adjacency_matrix(X, inverse_distance, time_step)
edge_index = adjacency_matrix.nonzero().t()
data = Data(x=X[time_step], y=y[time_step], edge_index=edge_index)
data_list.append(data)
return data_list
filenames = [f'data/task1_2/train/trajectory_{i}.npz' for i in range(900)]
split_point = 720 # take 20% of the data for validation
# Create two subsets for train/validation split from the filenames list
train_filenames = filenames[:split_point]
val_filenames = filenames[split_point:]
train_dataset = MyDataset(root='data/task1_2/train', filenames=train_filenames)
train_dataset = torch.utils.data.ConcatDataset(train_dataset)
train_dataloader = DataLoader(train_dataset, batch_size=32)
# Create validation dataset from val_filenames
val_dataset = MyDataset(root='data/task1_2/train', filenames=val_filenames)
val_dataset = torch.utils.data.ConcatDataset(val_dataset)
val_dataloader = DataLoader(val_dataset, batch_size=32)
# Prepare for validation data set
test_filenames = [f'data/task1_2/test/trajectory_{i}.npz' for i in range(901, 1000)]
test_dataset = MyDataset(root='data/task1_2/test', filenames=test_filenames)
test_dataset = torch.utils.data.ConcatDataset(test_dataset)
test_dataloader = DataLoader(test_dataset, batch_size=32)
this is the code for accessing the dataset and processing the data.
now i want to code the best model architecture for this data and task. can you come up with a very robust GNN which will perform very well for this tasl? it can be any type of GNNASSISTANT
Certainly! Given the task of predicting the energy of a system of atoms evolving in a 2D unit cell, we need a GNN model that captures the interactions between the atoms effectively. One robust architecture for such a problem is the Graph Isomorphism Network (GIN), which is known for its strong representational power.
Here's a robust GNN architecture using GIN layers, combined with other techniques like batch normalization and dropout to improve performance. We’ll also implement residual connections to make the model deeper and more powerful.
### Implementation
#### Step 1: Define the GNN Model
First, we define the `GINEnergyPredictor` model. This model will consist of multiple GIN layers, followed by batch normalization, activation functions, and dropout for regularization.
```python
import torch
from torch_geometric.nn import GINConv, global_add_pool
from torch.nn import Sequential, Linear, ReLU
class GINEnergyPredictor(torch.nn.Module):
def __init__(self, in_channels, hidden_channels, out_channels, num_layers, dropout):
super(GINEnergyPredictor, self).__init__()
self.conv_layers = torch.nn.ModuleList()
self.bns = torch.nn.ModuleList()
for i in range(num_layers):
if i == 0:
self.conv_layers.append(
GINConv(
Sequential(
Linear(in_channels, hidden_channels),
ReLU(),
Linear(hidden_channels, hidden_channels)
)
)
)
else:
self.conv_layers.append(
GINConv(
Sequential(
Linear(hidden_channels, hidden_channels),
ReLU(),
Linear(hidden_channels, hidden_channels)
)
)
)
self.bns.append(torch.nn.BatchNorm1d(hidden_channels))
self.final_lin = Sequential(
Linear(hidden_channels, hidden_channels),
ReLU(),
Linear(hidden_channels, out_channels)
)
self.dropout = torch.nn.Dropout(dropout)
def forward(self, data):
x, edge_index = data.x, data.edge_index
for conv, bn in zip(self.conv_layers, self.bns):
x = conv(x, edge_index)
x = bn(x)
x = F.relu(x)
x = self.dropout(x)
# Global pooling layer to aggregate node features into a graph feature
x = global_add_pool(x, data.batch)
x = self.final_lin(x)
return x
```
#### Step 2: Training Script
We train and validate the model with the processed data.
```python
import torch.nn.functional as F
from torch_geometric.data import DataLoader
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = GINEnergyPredictor(in_channels=5, hidden_channels=64, out_channels=1, num_layers=5, dropout=0.5).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = torch.nn.MSELoss()
def train(model, train_loader):
model.train()
total_loss = 0
for data in train_loader:
data = data.to(device)
optimizer.zero_grad()
output = model(data).squeeze()
loss = criterion(output, data.y.squeeze())
loss.backward()
optimizer.step()
total_loss += loss.item() * data.num_graphs
return total_loss / len(train_loader.dataset)
def validate(model, val_loader):
model.eval()
total_loss = 0
with torch.no_grad():
for data in val_loader:
data = data.to(device)
output = model(data).squeeze()
loss = criterion(output, data.y.squeeze())
total_loss += loss.item() * data.num_graphs
return total_loss / len(val_loader.dataset)
# Training loop
best_val_loss = float('inf')
patience = 10
num_epochs = 100
patience_counter = 0
for epoch in range(1, num_epochs + 1):
train_loss = train(model, train_dataloader)
val_loss = validate(model, val_dataloader)
if val_loss < best_val_loss:
best_val_loss = val_loss
torch.save(model.state_dict(), "best_model.pt")
patience_counter = 0
else:
patience_counter += 1
if patience_counter == patience:
print(f"Early stopping at epoch {epoch}")
break
print(f"Epoch {epoch}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}")
```
#### Step 3: Evaluate Model on Test Set
Finally, we load the best model and evaluate it on the test set.
```python
model.load_state_dict(torch.load("best_model.pt"))
model.eval()
def test(model, test_loader):
model.eval()
total_loss = 0
with torch.no_grad():
for data in test_loader:
data = data.to(device)
output = model(data).squeeze()
loss = criterion(output, data.y.squeeze())
total_loss += loss.item() * data.num_graphs
return total_loss / len(test_loader.dataset)
test_loss = test(model, test_dataloader)
print(f"Test Loss: {test_loss:.4f}")
```
### Summary
1. **Model Definition**: A robust GNN model with multiple GIN layers, batch normalization, ReLU activation, and dropout.
2. **Training**: Implemented a training loop with early stopping based on validation loss.
3. **Evaluation**: Evaluated the saved best model on the test set.
This model architecture and training approach should give you a strong baseline for predicting the energy of the atomic system accurately. Feel free to adjust hyperparameters like `num_layers`, `hidden_channels`, and `dropout` for fine-tuning the model's performance.