inverse_pinn.py

# -*- coding: utf-8 -*-
"""
Created on Sun Jun 18 16:33:15 2023

@author: Tomamso Giacometti
"""
#In this script the aim is to infer two parameters of the differential equation system: lambda and gamma.
#For the fit will be used data generated by scipy (with or without noise) only of the y1 and z variables.
import torch
import numpy as np
import plots
from scipy.integrate import odeint
from model import PINN_inverse
from utils import pde_scipy, inverse_pinn_data_gen

def main():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    #Data generation for the inverse problem
    #Parameters 
    lam = 0.2
    nu = 0.33
    gamma = 2.
    delta = 0.33
    params = np.array([lam,nu,gamma,delta])
    #I want to infer two parameters (lam and gamma), assuming already known the others (nu and delta)
    params_known = np.array([nu,delta])
    params_to_infer = np.array([1,1]) # Starting values for the parameters to infer lam and gamma 
    
    #Initial conditions
    x1 = 6
    x2 = 5
    y1 = 0
    z = 0
    init = np.array([x1,x2,y1,z])
    
    #Time domain
    ub_time = 10. #Upper bound for the time domain during the integration
    time = np.linspace(0, ub_time, 1000)
    
    #Solution of scipy
    y = odeint(pde_scipy, init, time, args=(params,))
    
    
    ##PINN solution
    # Create PINN instance
    torch.manual_seed(0)
    pinn = PINN_inverse().to(device)
    pinn.add_pde_parameters_known(params_known) # Add the parameters of the pde as buffers
    pinn.add_pde_parameters_to_infer(params_to_infer) # Add the parameters of the pde as torch.Parameters
    pinn.add_initial_condition(init) # Add the initial conditions, starting time of ic by default is zero
    
    # Convert time for pde loss to tensors
    # The input and target must be float numbers
    t = torch.from_numpy(time).float().view(-1, 1).to(device).requires_grad_(True)
    
    # Generating data in different time points
    np.random.seed(123)
    n_data = int(ub_time*2) # Number of data points to generate, I suppose, for now, that data is collected twice a day
    time_data_points = np.linspace(0, ub_time, n_data)
    data = inverse_pinn_data_gen(init, time_data_points, params, True, std=0.5) # Data with noise
    data = torch.from_numpy(data).float()
    
    # Set training parameters
    epochs = 10000
    optimizer = torch.optim.Adam(pinn.parameters(), lr=1e-3)
    loss_history = []
    
    #Training performed with Adam (or the one setted above)
    print(f'Adam training ({epochs} epochs):')
    for epoch in range(epochs):
        loss = pinn.train_step(t, data, optimizer).item() # Single train step defined in the PINN class
        loss_history.append(loss)
        if (epoch + 1) % (epochs / 20) == 0:
            print(f"{epoch/epochs*100:.1f}% -> Loss: ", loss)
            
    # Setting up the network for the LBFGS training
    optimizer, closure = pinn.set_up_lbfgs(t, data)
    
    # Training performed with LBFGS
    epochs = 500
    print(f'LBFGS training ({epochs} epochs):')
    for epoch in range(epochs):
        optimizer.step(closure)
        loss = closure()
        loss_history.append(loss.item())
        if torch.isnan(loss):
            error = 'The lbfgs training get a nan!\n'
            error += 'Especially with noisy data, lbfgs can produce errors due to convergence problems.\n'
            error += 'Possible solution can be modify lbfgs parameters (like history_size) in the model.py file.\n'
            error += 'If the error persist, the lbfgs training can be simply avoided, Adam training may be enough.'
            raise RuntimeError(error)
        if (epoch + 1) % (epochs / 20) == 0:
            print(f"{epoch/epochs*100:.1f}% -> Loss: ", loss.item())
            
    #Inferred parameters
    print(f'lambda = {pinn.params_to_infer[0].item():.5f}, gamma = {pinn.params_to_infer[1].item():.5f}')
    
    # Plots
    plots.plot_loss(loss_history)
    plots.plot_solution_pinn_inverse(pinn, time, data, y)
    
    
    

if __name__ == '__main__':
    main()