import torch
import torch.nn as nn
import torch.optim as optim

# Define the neural network model
class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(1, 10)  # Input layer to hidden layer
        self.relu = nn.ReLU()        # Activation function
        self.fc2 = nn.Linear(10, 1)  # Hidden layer to output layer

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# Instantiate the model
model = SimpleNN()

# Define the loss function and optimizer
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)

# Simulated training data
F_data = torch.rand(100, 1) * 10  # Force values between 0 and 10
A_data = F_data / 2               # a = F/m, with m = 2

# Training the network with physics-informed loss
lambda_reg = 1.0  # Regularization parameter

for epoch in range(1000):  # Training for 1000 epochs
    model.train()
    
    optimizer.zero_grad()
    
    # Forward pass
    A_pred = model(F_data)
    
    # Data-driven loss (mean squared error)
    loss_data = criterion(A_pred, A_data)
    
    # Physics-informed loss (F = 2 * a)
    physics_loss = criterion(F_data, 2 * A_pred)
    
    # Total loss
    loss = loss_data + lambda_reg * physics_loss
    
    # Backward pass and optimization
    loss.backward()
    optimizer.step()
    
    if epoch % 100 == 0:
        print(f'Epoch {epoch}, Loss: {loss.item()}')

# Testing the model
model.eval()
with torch.no_grad():
    test_F = torch.tensor([[5.0]])  # Example test input
    test_A_pred = model(test_F)
    print(f'Predicted acceleration for F=5: {test_A_pred.item()}')

Epoch 0, Loss: 33.23193359375
Epoch 100, Loss: 0.0010111266747117043
Epoch 200, Loss: 0.0002765950048342347
Epoch 300, Loss: 9.012725058710203e-05
Epoch 400, Loss: 1.9460174371488392e-05
Epoch 500, Loss: 3.5050049973506248e-06
Epoch 600, Loss: 1.844927510319394e-06
Epoch 700, Loss: 1.7233555809070822e-06
Epoch 800, Loss: 1.7036819599525188e-06
Epoch 900, Loss: 1.6891971199584077e-06
Predicted acceleration for F=5: 2.500002861022949