Fix bug in fhn and hopf OC (#255)

* fix bug in OC of fhn and hopf for diffusive coupling

* fix style fhn

* fix style hopf

* fix style fhn

* fix style hopf

neurolib/control/optimal_control/oc_fhn/oc_fhn.py

@@ -90,6 +90,9 @@ def compute_hx(self):
         """
         return compute_hx(
             self.model_params,
+            self.model.params["K_gl"],
+            self.model.params["Cmat"],
+            self.model.params["coupling"],
             self.N,
             self.dim_vars,
             self.T,
@@ -106,7 +109,6 @@ def compute_hx_nw(self):
         return compute_hx_nw(
             self.model.params["K_gl"],
             self.model.params["Cmat"],
-            self.model.params["coupling"],
             self.N,
             self.dim_vars,
             self.T,

neurolib/control/optimal_control/oc_hopf/oc_hopf.py

@@ -90,6 +90,9 @@ def compute_hx(self):
         """
         return compute_hx(
             self.model_params,
+            self.model.params["K_gl"],
+            self.model.params["Cmat"],
+            self.model.params["coupling"],
             self.N,
             self.dim_vars,
             self.T,
@@ -106,7 +109,6 @@ def compute_hx_nw(self):
         return compute_hx_nw(
             self.model.params["K_gl"],
             self.model.params["Cmat"],
-            self.model.params["coupling"],
             self.N,
             self.dim_vars,
             self.T,

neurolib/models/fhn/timeIntegration.py

@@ -287,6 +287,9 @@ def jacobian_fhn(
 @numba.njit
 def compute_hx(
     model_params,
+    K_gl,
+    cmat,
+    coupling,
     N,
     V,
     T,
@@ -295,8 +298,14 @@ def compute_hx(
 ):
     """Jacobians  of FHN model wrt. its 'state_vars' at each time step.
 
-    :param model_params:    Ordered tuple of parameters in the FHN Model in order
-    :type model_params:     tuple of float
+    :param model_params:        Ordered tuple of parameters in the FHN Model in order
+    :type model_params:         tuple of float
+    :param K_gl:                Model parameter of global coupling strength.
+    :type K_gl:                 float
+    :param cmat:                Model parameter, connectivity matrix.
+    :type cmat:                 ndarray
+    :param coupling:            Model parameter, which specifies the coupling type. E.g. "additive" or "diffusive".
+    :type coupling:             str
     :param N:                   Number of nodes in the network.
     :type N:                    int
     :param V:                   Number of system variables.
@@ -316,14 +325,17 @@ def compute_hx(
     for n in range(N):  # Iterate through nodes.
         for t in range(T):
             hx[n, t, :, :] = jacobian_fhn(model_params, dyn_vars[n, sv["x"], t], V, sv)
+
+            if coupling == "diffusive":
+                for l in range(N):
+                    hx[n, t, sv["x"], sv["x"]] += K_gl * cmat[n, l]
     return hx
 
 
 @numba.njit
 def compute_hx_nw(
     K_gl,
     cmat,
-    coupling,
     N,
     V,
     T,
@@ -335,8 +347,6 @@ def compute_hx_nw(
     :type K_gl:      float
     :param cmat:     Model parameter, connectivity matrix.
     :type cmat:      ndarray
-    :param coupling: Model parameter, which specifies the coupling type. E.g. "additive" or "diffusive".
-    :type coupling:  str
     :param N:        Number of nodes in the network.
     :type N:         int
     :param V:        Number of system variables.
@@ -353,9 +363,9 @@ def compute_hx_nw(
 
     for n1 in range(N):
         for n2 in range(N):
-            hx_nw[n1, n2, :, sv["x"], sv["x"]] = K_gl * cmat[n1, n2]  # term corresponding to additive coupling
-            if coupling == "diffusive":
-                hx_nw[n1, n1, :, sv["x"], sv["x"]] += -K_gl * cmat[n1, n2]
+            hx_nw[n1, n2, :, sv["x"], sv["x"]] = (
+                K_gl * cmat[n1, n2]
+            )  # term corresponding to both diffusive and additive coupling
 
     return -hx_nw
 

neurolib/models/hopf/timeIntegration.py

@@ -267,6 +267,9 @@ def jacobian_hopf(
 @numba.njit
 def compute_hx(
     model_params,
+    K_gl,
+    cmat,
+    coupling,
     N,
     V,
     T,
@@ -277,7 +280,13 @@ def compute_hx(
 
     :param model_params:    Ordered tuple of parameters in the Hopf Model in order
     :type model_params:     tuple of float
-    :param N:               Number of network nodes.
+    :param K_gl:            Model parameter of global coupling strength.
+    :type K_gl:             float
+    :param cmat:            Model parameter, connectivity matrix.
+    :type cmat:             ndarray
+    :param coupling:        Model parameter, which specifies the coupling type. E.g. "additive" or "diffusive".
+    :type coupling:         str
+    :param N:               Number of nodes in the network.
     :type N:                int
     :param V:               Number of state variables.
     :type V:                int
@@ -303,14 +312,17 @@ def compute_hx(
                 dyn_vars[n, sv["y"], t],
                 sv,
             )
+
+            if coupling == "diffusive":
+                for l in range(N):
+                    hx[n, t, sv["x"], sv["x"]] += K_gl * cmat[n, l]
     return hx
 
 
 @numba.njit
 def compute_hx_nw(
     K_gl,
     cmat,
-    coupling,
     N,
     V,
     T,
@@ -322,8 +334,6 @@ def compute_hx_nw(
     :type K_gl:         float
     :param cmat:        Model parameter, connectivity matrix.
     :type cmat:         ndarray
-    :param coupling:    Model parameter, which specifies the coupling type. E.g. "additive" or "diffusive".
-    :type coupling:     str
     :param N:           Number of nodes in the network.
     :type N:            int
     :param V:           Number of system variables.
@@ -340,9 +350,9 @@ def compute_hx_nw(
 
     for n1 in range(N):
         for n2 in range(N):
-            hx_nw[n1, n2, :, sv["x"], sv["x"]] = K_gl * cmat[n1, n2]  # term corresponding to additive coupling
-            if coupling == "diffusive":
-                hx_nw[n1, n1, :, sv["x"], sv["x"]] += -K_gl * cmat[n1, n2]
+            hx_nw[n1, n2, :, sv["x"], sv["x"]] = (
+                K_gl * cmat[n1, n2]
+            )  # corresponding to both diffusive and additive coupling
 
     return -hx_nw