Fixed and improved BundleSolver1D

[at-chantada]

There are two modifications to the use of BundleSolver1D and BundleIVP in this pull request.

The first is a fix of BundleIVP, where it didn't work in the case of a bundle that included conditions for different variables in a system of ODEs. To make the fix, now the bundle_conditions input takes a dict, where in each entry the key represents the condition, and the corresponding value, the index in the tuple used for theta_min and theta_max in BundleSolver1D, where the ranges are indicated for that condition.

The second is an improvement, where now it is not necessary to put the initial conditions as inputs in your ode_system, when including them in the bundle.

To illustrate the changes here is an example where a network is trained to be a solution for the damped harmonic oscillator, with different values for the initial time t_0 (between -1 and 0) initial position x1_0 (between 0 and 1) and the initial velocity x2_0 (between 1 and 2):

First let's create the analytical solution to check how good is the network once it's trained:

import numpy as np

def damp_harm_sol(t, t_0, x_0, x_0_prime, m, b, k):
    x_t = 0

    # Under-damping:
    if b**2 < 4*m*k:

        # Frequency:
        w = np.sqrt(np.abs((b**2)-4*m*k))/(2*m)

        # Exponent factor:
        e = b/(2*m)

        # Parameters:
        c_1 = x_0
        c_2 = (x_0_prime+e*x_0)/w

        # Function:
        x = np.exp(-e*(t-t_0))*(c_1*np.cos(w*(t-t_0))+c_2*np.sin(w*(t-t_0)))
        x_prime = -e*np.exp(-e*(t-t_0))*(c_1*np.cos(w*(t-t_0))+c_2*np.sin(w*(t-t_0))) + w*np.exp(-e*(t-t_0))*(-c_1*np.sin(w*(t-t_0))+c_2*np.cos(w*(t-t_0)))

    return x, x_prime

The code to train the network:

import matplotlib.pyplot as plt
from neurodiffeq import diff      # the differentiation operation
from neurodiffeq.conditions import BundleIVP   # the initial condition
from neurodiffeq.solvers import BundleSolver1D  # the solver

# Damped oscilator paramters
m = 1
b = 1
k = 2
x1_0_min = 0
x1_0_max = 1
x2_0_min = 1
x2_0_max = 2
t_0_min = -1
t_0_max = 0
t_f = np.pi

# specify the system of ODEs
def parametric_damp (x1, x2, t):
    A = diff(x1, t) - x2
    B = m*diff(x2, t) + b*x2 + k*x1   

    return A, B

# specify the initial conditions

init_vals_damp = [BundleIVP(bundle_conditions={'t_0': 0, 'u_0': 1}), 
                  BundleIVP(bundle_conditions={'t_0': 0, 'u_0': 2})]  
# Specify in bundle_conditions, the conditions that are to be included in the bundle.
# Where the values in the dict asociated with each condition
# correspond to the index in the tuple in theta_min and theta_max.

# solve the ODE
solution_damp = BundleSolver1D(
    ode_system=parametric_damp, conditions=init_vals_damp, t_min=t_0_min, t_max=t_f,
    theta_min=(t_0_min, x1_0_min, x2_0_min),  # The index of the tuple must match reflect correctly with bundle_conditions, if conditions are part of the bundle
    theta_max=(t_0_max, x1_0_max, x2_0_max)  # The index of the tuple must match reflect correctly with bundle_conditions, if conditions are part of the bundle
    )
solution_damp.fit(max_epochs=5000)
solution = solution_damp.get_solution()

Finally some plots:

for t_0 in [-0.9, -0.1]:
  ts = np.linspace(t_0, t_f, 50)
  for x1_0, x2_0 in zip([0.1, 0.1, 0.9, 0.9] , [1.1, 1.9, 1.1, 1.9]):
    x1_net, x2_net = solution(ts, t_0 * np.ones(len(ts)), x1_0 * np.ones(len(ts)), x2_0 * np.ones(len(ts)), to_numpy=True)
    x1_ana, x2_ana = damp_harm_sol(ts, t_0, x1_0, x2_0, m, b, k)

    fig, axs = plt.subplots(1, 1)

    axs.plot(ts, x1_net, label='ANN-based solution of x1', color='C1')
    axs.plot(ts, x1_ana, '.',label='analytical solution of x1', color='C0')
    axs.plot(ts, x2_net, label='ANN-based solution of x2', color='C3')
    axs.plot(ts, x2_ana, '.',label='analytical solution of x2', color='C2')
    axs.set_xlabel('t')
    axs.set_ylabel('x')

    title = 'Damped Harmonic Oscilator \n (' + r'$m={},\;k={},\;b={},\;t_0={},\;x_0={},\;x^\prime_0={}$'.format(m, k, b, t_0, x1_0, x2_0)  + ')'
    axs.set_title(title)
    axs.legend()
    filename = 't={}_x={}_x_prime={}.png'.format(t_0, x1_0, x2_0)
    plt.savefig(filename)

fig, axs = plt.subplots(1, 1)

loss_damp = solution_damp.metrics_history

axs.plot(loss_damp['train_loss'], label='training loss')
axs.plot(loss_damp['valid_loss'], label='validation loss')
axs.set_xlabel('epochs')
axs.set_yscale('log')
axs.set_title('Loss during training')
axs.legend()
plt.savefig('loss.png')

[codecov-commenter]

Codecov Report

Merging #131 (4c2858d) into master (b400aff) will decrease coverage by 0.39%. The diff coverage is 16.66%.

@@            Coverage Diff             @@
##           master     #131      +/-   ##
==========================================
- Coverage   89.93%   89.53%   -0.40%     
==========================================
  Files          18       18              
  Lines        2940     2953      +13     
==========================================
  Hits         2644     2644              
- Misses        296      309      +13     

Impacted Files	Coverage Δ
neurodiffeq/solvers.py	84.53% <0.00%> (-3.04%)	:arrow_down:
neurodiffeq/conditions.py	93.87% <100.00%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more Δ = absolute <relative> (impact), ø = not affected, ? = missing data Powered by Codecov. Last update b400aff...4c2858d. Read the comment docs.

[at-chantada]

As always, thank you for the great work with the library @shuheng-liu. With respect to the names, there is no particular need for those names really. In the case of r is just what I came up in the moment to use for a generic vector of length N. And in the case of theta I followed the notation used on the paper that i mentioned in #118. But of course those names can be changed to be more intuitive, and avoid causing confusion in the future.