ENH: add implcit neural ode solver for gr4 model

[nghi-truyen]

• Implementation of implicit neural ode solver for gr4 model structure (backward of MLP in Fortran is needed).
• Optimize calculation for forward and backward of parameterization NN using Fortran element-wise operations instead of loop.
• Test of original structures vs NN-based structures
• For the unit test, only results on neural ode solver have been modified due to the change from explicit Euler to implicit Euler.
• Upgrade dev env to Python 3.12

Python script to test the forward and backward MLP in Fortran:

import smash
import numpy as np
# import also the forward and backward MLP that are wrapped to Python from Fortran solver

# Define weights and biases
x = np.random.uniform(size=5)
w1 = np.random.uniform(size=(3,5))
b1 = np.random.uniform(size=3)
w2 = np.random.uniform(size=(2,3))
b2 = np.random.uniform(size=2)
w3 = np.random.uniform(size=(4,2))
b3 = np.random.uniform(size=4)
# x =  [0.75189605 0.72820489 0.20674587 0.83382761 0.26712038]
# w1 =  [[0.92751443 0.36317825 0.19527185 0.69294418 0.99065735]
#  [0.86592888 0.58032284 0.11218693 0.22154433 0.42089211]
#  [0.19136336 0.42783286 0.85232179 0.05576818 0.91000867]]
# b1 =  [0.35388026 0.67698752 0.65528754]
# w2 =  [[0.39774009 0.03414731 0.18890095]
# [0.39704593 0.09149761 0.14763046]]
# b2 =  [0.5659279 0.0575433]
# w3 =  [[0.84487334 0.62280035]
#  [0.44611591 0.86117966]
#  [0.89682374 0.90132414]
#  [0.48648652 0.58286809]]
# b3 =  [0.20367196 0.0022146  0.30377352 0.74705191]

## Forward and backward in Python
#=============================
def forward_nn(x,w1,b1,w2,b2,w3,b3):
    f1 = np.dot(w1, x) + b1
    f1 = np.maximum(0.01*f1, f1)
    f2 = np.dot(w2, f1) + b2
    f2 = np.maximum(0.01*f2, f2)
    f3 = np.dot(w3, f2) + b3
    return np.tanh(f3)

def tanh_derivative(z):
    return 1 - np.tanh(z)**2

def leaky_relu_derivative(z):
    return np.where(z > 0, 1, 0.01)

def forward_and_intermediate_values(x, w1, b1, w2, b2, w3, b3):
    z1 = np.dot(w1, x) + b1
    f1 = np.maximum(0.01 * z1, z1)  # Leaky ReLU
    z2 = np.dot(w2, f1) + b2
    f2 = np.maximum(0.01 * z2, z2)  # Leaky ReLU
    z3 = np.dot(w3, f2) + b3
    f3 = np.tanh(z3)  # TanH
    return z1, f1, z2, f2, z3, f3

def compute_gradient(x, w1, b1, w2, b2, w3, b3):
    z1, f1, z2, f2, z3, f3 = forward_and_intermediate_values(x, w1, b1, w2, b2, w3, b3)

    jaco = np.zeros((4, 5))  # 4 outputs, 5 inputs
    
    for i in range(4):
        # Initial gradient from tanh
        grad = np.zeros(4)
        grad[i] = tanh_derivative(z3[i])
        
        # Layer 3 to 2
        grad = w3.T @ (grad * 1.0)
        grad = grad * leaky_relu_derivative(z2)
        
        # Layer 2 to 1
        grad = w2.T @ grad
        grad = grad * leaky_relu_derivative(z1)
        
        # Layer 1 to input
        grad = w1.T @ grad
       
        jaco[i, :] = grad
        
    return jaco

## Compare results with Fortran wrapped routines
# ========================================
forward_nn(x,w1,b1,w2,b2,w3,b3)
# array([0.98845486, 0.96215205, 0.99629925, 0.98416929])

wrap_forward_fortran(w1,b1,w2,b2,w3,b3,x)
# array([0.9884549 , 0.96215206, 0.99629927, 0.9841693 ], dtype=float32)

compute_gradient(x,w1,b1,w2,b2,w3,b3)
# array([[0.01523199, 0.00847754, 0.00775786, 0.01003801, 0.01935052],
#        [0.04482528, 0.02477705, 0.02169637, 0.02914938, 0.05543567],
#        [0.00604756, 0.00335816, 0.00302921, 0.00396785, 0.00761496],
#        [0.01535166, 0.00851368, 0.0076168 , 0.01004723, 0.01923354]])

wrap_backward_fortran(w1,b1,w2,b2,w3,b3,x)
# array([[0.01523196, 0.00847752, 0.00775785, 0.01003799, 0.01935049],
#        [0.04482526, 0.02477704, 0.02169636, 0.02914937, 0.05543565],
#        [0.00604753, 0.00335815, 0.00302919, 0.00396783, 0.00761493],
#        [0.01535165, 0.00851368, 0.00761679, 0.01004723, 0.01923352]],
#       dtype=float32)

[pag13]

C'est excellent Truyen, merci, impatient de voir les tests numeriques étendus de ces solveurs ;)

[nghi-truyen]

Merci bcp @pag13 pour ta review ;)

[asjeb]

Tout bon pour moi