From d3581d6c9a93a5994becd1367d93d3c80fbd1578 Mon Sep 17 00:00:00 2001
From: Ngo Nghi Truyen Huynh <129378719+nghi-truyen@users.noreply.github.com>
Date: Mon, 9 Sep 2024 13:51:28 +0200
Subject: [PATCH] FIX PR: apply suggestion changes from FC review regarding to
 hybird model branch (#311)

---
 smash/fcore/forward/forward_db.f90            | 19 ++++++++++++------
 smash/fcore/forward/forward_openmp_db.f90     | 19 ++++++++++++------
 .../routine/mwd_parameters_manipulation.f90   | 20 +++++++++++++------
 3 files changed, 40 insertions(+), 18 deletions(-)

diff --git a/smash/fcore/forward/forward_db.f90 b/smash/fcore/forward/forward_db.f90
index 422c0e36..5d849ff9 100644
--- a/smash/fcore/forward/forward_db.f90
+++ b/smash/fcore/forward/forward_db.f90
@@ -6377,12 +6377,15 @@ SUBROUTINE NN_PARAMETERS_GET_CONTROL_SIZE(setup, options, n)
     TYPE(SETUPDT), INTENT(IN) :: setup
     TYPE(OPTIONSDT), INTENT(IN) :: options
     INTEGER, INTENT(INOUT) :: n
-    n = options%optimize%nn_parameters(1)*setup%neurons(2)*setup%neurons&
-&     (1) + options%optimize%nn_parameters(2)*setup%neurons(2) + options&
-&     %optimize%nn_parameters(3)*setup%neurons(3)*setup%neurons(2) + &
-&     options%optimize%nn_parameters(4)*setup%neurons(3) + options%&
-&     optimize%nn_parameters(5)*setup%neurons(4)*setup%neurons(3) + &
-&     options%optimize%nn_parameters(6)*setup%neurons(4)
+    INTEGER :: i, n_w, n_b
+    INTRINSIC SIZE
+    n = 0
+    DO i=0,SIZE(setup%hidden_neuron)
+      n_w = options%optimize%nn_parameters(2*i+1)*setup%neurons(i+2)*&
+&       setup%neurons(i+1)
+      n_b = options%optimize%nn_parameters(2*i+2)*setup%neurons(i+2)
+      n = n + n_w + n_b
+    END DO
   END SUBROUTINE NN_PARAMETERS_GET_CONTROL_SIZE
 
   SUBROUTINE GET_CONTROL_SIZES(setup, mesh, options, nbk)
@@ -6793,6 +6796,7 @@ SUBROUTINE NN_PARAMETERS_FILL_CONTROL(setup, options, parameters)
     CHARACTER(len=lchar) :: name
     INTEGER :: j, k, l
     INTRINSIC SUM
+! NN parameters is fifth control kind
     j = SUM(parameters%control%nbk(1:4))
     IF (options%optimize%nn_parameters(1) .EQ. 1) THEN
       DO k=1,setup%neurons(1)
@@ -8248,6 +8252,7 @@ SUBROUTINE NN_PARAMETERS_FILL_PARAMETERS_D(setup, options, parameters&
     TYPE(PARAMETERSDT), INTENT(INOUT) :: parameters_d
     INTEGER :: j, k, l
     INTRINSIC SUM
+! NN parameters is fifth control kind
     j = SUM(parameters%control%nbk(1:4))
     IF (options%optimize%nn_parameters(1) .EQ. 1) THEN
       parameters_d%nn_parameters%weight_1 = 0.0_4
@@ -8343,6 +8348,7 @@ SUBROUTINE NN_PARAMETERS_FILL_PARAMETERS_B(setup, options, parameters&
     INTEGER :: j, k, l
     INTRINSIC SUM
     INTEGER :: branch
+! NN parameters is fifth control kind
     j = SUM(parameters%control%nbk(1:4))
     IF (options%optimize%nn_parameters(1) .EQ. 1) THEN
       DO k=1,setup%neurons(1)
@@ -8467,6 +8473,7 @@ SUBROUTINE NN_PARAMETERS_FILL_PARAMETERS(setup, options, parameters)
     TYPE(PARAMETERSDT), INTENT(INOUT) :: parameters
     INTEGER :: j, k, l
     INTRINSIC SUM
+! NN parameters is fifth control kind
     j = SUM(parameters%control%nbk(1:4))
     IF (options%optimize%nn_parameters(1) .EQ. 1) THEN
       DO k=1,setup%neurons(1)
diff --git a/smash/fcore/forward/forward_openmp_db.f90 b/smash/fcore/forward/forward_openmp_db.f90
index e97fdb2b..0c0d206c 100644
--- a/smash/fcore/forward/forward_openmp_db.f90
+++ b/smash/fcore/forward/forward_openmp_db.f90
@@ -6377,12 +6377,15 @@ SUBROUTINE NN_PARAMETERS_GET_CONTROL_SIZE(setup, options, n)
     TYPE(SETUPDT), INTENT(IN) :: setup
     TYPE(OPTIONSDT), INTENT(IN) :: options
     INTEGER, INTENT(INOUT) :: n
-    n = options%optimize%nn_parameters(1)*setup%neurons(2)*setup%neurons&
-&     (1) + options%optimize%nn_parameters(2)*setup%neurons(2) + options&
-&     %optimize%nn_parameters(3)*setup%neurons(3)*setup%neurons(2) + &
-&     options%optimize%nn_parameters(4)*setup%neurons(3) + options%&
-&     optimize%nn_parameters(5)*setup%neurons(4)*setup%neurons(3) + &
-&     options%optimize%nn_parameters(6)*setup%neurons(4)
+    INTEGER :: i, n_w, n_b
+    INTRINSIC SIZE
+    n = 0
+    DO i=0,SIZE(setup%hidden_neuron)
+      n_w = options%optimize%nn_parameters(2*i+1)*setup%neurons(i+2)*&
+&       setup%neurons(i+1)
+      n_b = options%optimize%nn_parameters(2*i+2)*setup%neurons(i+2)
+      n = n + n_w + n_b
+    END DO
   END SUBROUTINE NN_PARAMETERS_GET_CONTROL_SIZE
 
   SUBROUTINE GET_CONTROL_SIZES(setup, mesh, options, nbk)
@@ -6793,6 +6796,7 @@ SUBROUTINE NN_PARAMETERS_FILL_CONTROL(setup, options, parameters)
     CHARACTER(len=lchar) :: name
     INTEGER :: j, k, l
     INTRINSIC SUM
+! NN parameters is fifth control kind
     j = SUM(parameters%control%nbk(1:4))
     IF (options%optimize%nn_parameters(1) .EQ. 1) THEN
       DO k=1,setup%neurons(1)
@@ -8250,6 +8254,7 @@ SUBROUTINE NN_PARAMETERS_FILL_PARAMETERS_D(setup, options, parameters&
     TYPE(PARAMETERSDT), INTENT(INOUT) :: parameters_d
     INTEGER :: j, k, l
     INTRINSIC SUM
+! NN parameters is fifth control kind
     j = SUM(parameters%control%nbk(1:4))
     IF (options%optimize%nn_parameters(1) .EQ. 1) THEN
       parameters_d%nn_parameters%weight_1 = 0.0_4
@@ -8345,6 +8350,7 @@ SUBROUTINE NN_PARAMETERS_FILL_PARAMETERS_B(setup, options, parameters&
     INTEGER :: j, k, l
     INTRINSIC SUM
     INTEGER :: branch
+! NN parameters is fifth control kind
     j = SUM(parameters%control%nbk(1:4))
     IF (options%optimize%nn_parameters(1) .EQ. 1) THEN
       DO k=1,setup%neurons(1)
@@ -8469,6 +8475,7 @@ SUBROUTINE NN_PARAMETERS_FILL_PARAMETERS(setup, options, parameters)
     TYPE(PARAMETERSDT), INTENT(INOUT) :: parameters
     INTEGER :: j, k, l
     INTRINSIC SUM
+! NN parameters is fifth control kind
     j = SUM(parameters%control%nbk(1:4))
     IF (options%optimize%nn_parameters(1) .EQ. 1) THEN
       DO k=1,setup%neurons(1)
diff --git a/smash/fcore/routine/mwd_parameters_manipulation.f90 b/smash/fcore/routine/mwd_parameters_manipulation.f90
index f7b2b235..077bbb2f 100644
--- a/smash/fcore/routine/mwd_parameters_manipulation.f90
+++ b/smash/fcore/routine/mwd_parameters_manipulation.f90
@@ -748,12 +748,18 @@ subroutine nn_parameters_get_control_size(setup, options, n)
         type(OptionsDT), intent(in) :: options
         integer, intent(inout) :: n
 
-        n = options%optimize%nn_parameters(1)*setup%neurons(2)*setup%neurons(1) &
-        & + options%optimize%nn_parameters(2)*setup%neurons(2) &
-        & + options%optimize%nn_parameters(3)*setup%neurons(3)*setup%neurons(2) &
-        & + options%optimize%nn_parameters(4)*setup%neurons(3) &
-        & + options%optimize%nn_parameters(5)*setup%neurons(4)*setup%neurons(3) &
-        & + options%optimize%nn_parameters(6)*setup%neurons(4)
+        integer :: i, n_w, n_b
+
+        n = 0
+
+        do i = 0, size(setup%hidden_neuron)
+
+            n_w = options%optimize%nn_parameters(2*i + 1)*setup%neurons(i + 2)*setup%neurons(i + 1)
+            n_b = options%optimize%nn_parameters(2*i + 2)*setup%neurons(i + 2)
+
+            n = n + n_w + n_b
+
+        end do
 
     end subroutine nn_parameters_get_control_size
 
@@ -1241,6 +1247,7 @@ subroutine nn_parameters_fill_control(setup, options, parameters)
         character(lchar) :: name
         integer :: j, k, l
 
+        ! NN parameters is fifth control kind
         j = sum(parameters%control%nbk(1:4))
 
         if (options%optimize%nn_parameters(1) .eq. 1) then
@@ -1780,6 +1787,7 @@ subroutine nn_parameters_fill_parameters(setup, options, parameters)
 
         integer :: j, k, l
 
+        ! NN parameters is fifth control kind
         j = sum(parameters%control%nbk(1:4))
 
         if (options%optimize%nn_parameters(1) .eq. 1) then