Add option to return network, example generator, optimizer and loss function
@@ -92,7 +92,7 @@ def solve(ode, condition, t_min, t_max,
net=None, example_generator=None, shuffle=True,
optimizer=None, criterion=None, batch_size=16,
max_epochs=100000, tol=1e-4,
- monitor=None):
+ monitor=None, return_internal=False):
"""
Train a neural network to solve an ODE.
@@ -115,13 +115,18 @@ def solve(ode, condition, t_min, t_max,
:param monitor: a Monitor instance
"""
nets = None if not net else [net]
- solution, loss_history = solve_system(
+ returned_tuple = solve_system(
ode_system=lambda x, t: [ode(x, t)], conditions=[condition],
t_min=t_min, t_max=t_max, nets=nets,
example_generator=example_generator, shuffle=shuffle,
optimizer=optimizer, criterion=criterion, batch_size=batch_size,
- max_epochs=max_epochs, tol=tol, monitor=monitor
+ max_epochs=max_epochs, tol=tol, monitor=monitor, return_internal=return_internal
)
+ if return_internal:
+ solution, loss_history, internal = returned_tuple
+ return lambda t: solution(t)[0], loss_history, internal
+ else:
+ solution, loss_history = returned_tuple
return lambda t: solution(t)[0], loss_history
@@ -129,7 +134,7 @@ def solve_system(ode_system, conditions, t_min, t_max,
nets=None, example_generator=None, shuffle=True,
optimizer=None, criterion=None, batch_size=16,
max_epochs=100000, tol=1e-4,
- monitor=None):
+ monitor=None, return_internal=False):
"""
Train a neural network to solve an ODE.
@@ -170,6 +175,15 @@ def solve_system(ode_system, conditions, t_min, t_max,
if not criterion:
criterion = nn.MSELoss()
+ if return_internal:
+ internal = {
+ 'nets': nets,
+ 'conditions': conditions,
+ 'example_generator': example_generator,
+ 'optimizer': optimizer,
+ 'criterion': criterion
+ }
+
n_examples = example_generator.size
zeros = torch.zeros(batch_size)
@@ -227,4 +241,7 @@ def solve_system(ode_system, conditions, t_min, t_max,
if loss_history[-1] > tol:
print('The solution has not converged.')
+ if return_internal:
+ return solution, loss_history, internal
+ else:
return solution, loss_history
