Spaces:
Sleeping
Sleeping
File size: 3,107 Bytes
723b602 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | # Conditional import for neural ODE solver to handle cases where torch is not available
try:
import torch
import torch.nn as nn
import torch.optim as optim
TORCH_AVAILABLE = True
except ImportError:
TORCH_AVAILABLE = False
import numpy as np
if TORCH_AVAILABLE:
class NeuralODE(nn.Module):
"""
Neural Network surrogate for solving ODEs
Takes time t as input and outputs (x(t), y(t))
"""
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Linear(1, 64),
nn.Tanh(),
nn.Linear(64, 64),
nn.Tanh(),
nn.Linear(64, 64),
nn.Tanh(),
nn.Linear(64, 2), # Output x, y
)
def forward(self, t):
return self.net(t)
def train_neural_ode(t_train, x_train, y_train, epochs=2000, lr=1e-3):
"""
Train a neural network to approximate the solution on the training interval
Args:
t_train: array of time points for training
x_train: array of x values for training
y_train: array of y values for training
epochs: number of training epochs
lr: learning rate
Returns:
Trained model
"""
# Prepare training data
t_train_tensor = torch.tensor(t_train[:, None], dtype=torch.float32)
xy_train_tensor = torch.tensor(np.column_stack([x_train, y_train]), dtype=torch.float32)
# Initialize model
model = NeuralODE()
optimizer = optim.Adam(model.parameters(), lr=lr)
loss_fn = nn.MSELoss()
# Training loop
for epoch in range(epochs):
optimizer.zero_grad()
pred = model(t_train_tensor)
loss = loss_fn(pred, xy_train_tensor)
loss.backward()
optimizer.step()
if epoch % 500 == 0:
print(f"Epoch {epoch}, Loss: {loss.item():.6f}")
return model
def predict_with_neural_ode(model, t_eval):
"""
Predict solution using the trained neural network
Args:
model: trained NeuralODE model
t_eval: array of time points to evaluate
Returns:
x_pred, y_pred arrays
"""
if model is None:
return None, None
with torch.no_grad():
t_tensor = torch.tensor(t_eval[:, None], dtype=torch.float32)
pred = model(t_tensor).numpy()
x_pred = pred[:, 0]
y_pred = pred[:, 1]
return x_pred, y_pred
else:
def train_neural_ode(t_train, x_train, y_train, epochs=2000, lr=1e-3):
"""
Placeholder function when torch is not available
"""
print("PyTorch not available, skipping neural network training")
return None
def predict_with_neural_ode(model, t_eval):
"""
Placeholder function when torch is not available
"""
return None, None |