Optimal Control extension

README.md

@@ -95,6 +95,10 @@ neurolib/	 				# Main module
 ├── optimize/ 					# Optimization submodule
 	├── evolution/ 				# Evolutionary optimization
 	└── exploration/ 			# Parameter exploration
+├── control/optimal_control/			# (Optimal) control submodule
+	├── oc.py 				# Optimal control base class
+	├── cost_functions.py 			# cost functions for OC
+	├── /.../ 				# Implemented OC models
 ├── data/ 					# Empirical datasets (structural, functional)
 ├── utils/					# Utility belt
 	├── atlases.py				# Atlases (Region names, coordinates)
@@ -120,6 +124,7 @@ Example [IPython Notebooks](examples/) on how to use the library can be found in
 - [Example 0.6](https://mybinder.org/v2/gh/neurolib-dev/neurolib/master?filepath=examples%2Fexample-0.6-custom-model.ipynb) - Minimal example of how to implement your own model in `neurolib`
 - [Example 1.2](https://mybinder.org/v2/gh/neurolib-dev/neurolib/master?filepath=examples%2Fexample-1.2-brain-network-exploration.ipynb) - Parameter exploration of a brain network and fitting to BOLD data
 - [Example 2.0](https://mybinder.org/v2/gh/neurolib-dev/neurolib/master?filepath=examples%2Fexample-2-evolutionary-optimization-minimal.ipynb) - A simple example of the evolutionary optimization framework 
+- [Example 5.2](https://mybinder.org/v2/gh/neurolib-dev/neurolib/master?filepath=examples%2Fexample-5.2-oc-wc-model-deterministic.ipynb) - Example of optimal control of the noise-free Wilson-Cowan model
 
 A basic overview of the functionality of `neurolib` is also given in the following. 
 
@@ -282,6 +287,31 @@ This will gives us a summary of the last generation and plots a distribution of
   <img src="https://github.com/neurolib-dev/neurolib/raw/master/resources/evolution_animated.gif">
 </p>
 
+### Optimal control
+The optimal control module enables to compute efficient stimulation for your neural model. If you know how your output should look like, this module computes the optimal input. Detailes example notebooks can be found in the [example folder](https://github.com/neurolib-dev/neurolib/tree/master/examples) (examples 5.1, 5.2, 5.3, 5.4). In optimal control computations, you trade precision with respect to a target against control strength. You can determine how much each contribution affects the results, by setting weights accordingly.
+
+To compute an optimal control signal, you need to create a model (e.g., an FHN model) and define a target state (e.g., a sine curve with period 2).
+```python
+from neurolib.models.fhn import FHNModel
+model = FHNModel()
+
+duration = 10.
+model.params["duration"] = duration
+dt = model.params["dt"]
+
+period = 2.
+target = np.sin(2. * np.pi * np.arange(0, duration+dt, dt) / period)
+```
+You can then create a controlled model and run the iterative optimization to find the most efficient control input. The optimal control and the controlled model activity can be taken from the controlled model.
+```python
+model_controlled = oc_fhn.OcFhn(model, target)
+model_controlled.optimize(500) # run 500 iterations
+optimal_control = model_controlled.control
+optimal_state = model_controlled.get_xs()
+```
+
+For a comprehensive study on optimal control of the Wilson-Cowan model based on the neurolib optimal control module, see Salfenmoser, L. & Obermayer, K. Optimal control of a Wilson–Cowan model of neural population dynamics. Chaos 33, 043135 (2023). https://doi.org/10.1063/5.0144682.
+
 ## More information
 
 ### Built With
@@ -313,9 +343,11 @@ Cakan, C., Jajcay, N. & Obermayer, K. neurolib: A Simulation Framework for Whole
 
 Caglar Cakan (cakan@ni.tu-berlin.de)  
 Department of Software Engineering and Theoretical Computer Science, Technische Universität Berlin, Germany  
-Bernstein Center for Computational Neuroscience Berlin, Germany  
+Bernstein Center for Computational Neuroscience Berlin, Germany 
 
 ### Acknowledgments
 This work was supported by the Deutsche Forschungsgemeinschaft (DFG, German Research Foundation) with the project number 327654276 (SFB 1315) and the Research Training Group GRK1589/2.
 
+The optimal control module was developed by Lena Salfenmoser and Martin Krück supported by the DFG project 163436311 (SFB 910).
+
 <!--end-include-in-documentation-->

neurolib/control/optimal_control/cost_functions.py

@@ -0,0 +1,210 @@
+import numpy as np
+import numba
+
+
+@numba.njit
+def accuracy_cost(x, target_timeseries, weights, cost_matrix, dt, interval=(0, None)):
+    """Total cost related to the accuracy, weighted sum of contributions.
+
+    :param x:               State of dynamical system.
+    :type x:                np.ndarray
+    :param target_timeseries:    Target state.
+    :type target_timeseries:     np.darray
+    :param weights:         Dictionary of weights.
+    :type weights:          dictionary
+    :param cost_matrix:     Matrix of channels to take into account
+    :type cost_matrix:      ndarray
+    :param dt:              Time step.
+    :type dt:               float
+    :param interval:        (t_start, t_end). Indices of start and end point of the slice (both inclusive) in time
+                            dimension. Only 'int' positive index-notation allowed (i.e. no negative indices or 'None').
+    :type interval:         tuple, optional
+
+    :return:                Accuracy cost.
+    :rtype:                 float
+    """
+
+    cost_timeseries = np.zeros((target_timeseries.shape))
+
+    # timeseries of control vector is weighted sum of contributing cost functionals
+    if weights["w_p"] != 0.0:
+        cost_timeseries += weights["w_p"] * precision_cost(x, target_timeseries, cost_matrix, interval)
+
+    cost = 0.0
+    # integrate over nodes, channels, and time
+    if weights["w_p"] != 0.0:
+        for n in range(x.shape[0]):
+            for v in range(x.shape[1]):
+                for t in range(interval[0], interval[1]):
+                    cost += cost_timeseries[n, v, t] * dt
+
+    return cost
+
+
+@numba.njit
+def derivative_accuracy_cost(x, target_timeseries, weights, cost_matrix, interval=(0, None)):
+    """Derivative of the 'accuracy_cost' wrt. the state 'x'.
+
+    :param x:               State of dynamical system.
+    :type x:                np.ndarray
+    :param target_timeseries:    Target state.
+    :type target_timeseries:     np.darray
+    :param weights:         Dictionary of weights.
+    :type weights:          dictionary
+    :param cost_matrix:     Matrix of channels to take into account
+    :type cost_matrix:      ndarray
+    :param interval:        (t_start, t_end). Indices of start and end point of the slice (both inclusive) in time
+                            dimension. Only 'int' positive index-notation allowed (i.e. no negative indices or 'None').
+    :type interval:         tuple, optional
+
+    :return:                Accuracy cost derivative.
+    :rtype:                 ndarray
+    """
+
+    der = np.zeros((target_timeseries.shape))
+
+    if weights["w_p"] != 0.0:
+        der += weights["w_p"] * derivative_precision_cost(x, target_timeseries, cost_matrix, interval)
+
+    return der
+
+
+@numba.njit
+def precision_cost(x_sim, x_target, cost_matrix, interval=(0, None)):
+    """Summed squared difference between target and simulation within specified time interval weighted by w_p.
+       Penalizes deviation from the target.
+
+    :param x_sim:       N x V x T array that contains the simulated time series.
+    :type x_sim:        np.ndarray
+    :param x_target:    N x V x T array that contains the target time series.
+    :type x_target:     np.ndarray
+    :param cost_matrix: N x V binary matrix that defines nodes and channels of precision measurement. Defaults to
+                             None.
+    :type cost_matrix:  np.ndarray
+    :param interval:    (t_start, t_end). Indices of start and end point of the slice (both inclusive) in time
+                        dimension. Only 'int' positive index-notation allowed (i.e. no negative indices or 'None').
+    :type interval:     tuple
+
+    :return:            Precision cost for time interval.
+    :rtype:             float
+    """
+
+    cost = np.zeros((x_target.shape))
+
+    # integrate over nodes, channels, and time
+    for n in range(x_target.shape[0]):
+        for v in range(x_target.shape[1]):
+            for t in range(interval[0], interval[1]):
+                cost[n, v, t] = 0.5 * cost_matrix[n, v] * (x_target[n, v, t] - x_sim[n, v, t]) ** 2
+
+    return cost
+
+
+@numba.njit
+def derivative_precision_cost(x_sim, x_target, cost_matrix, interval):
+    """Derivative of 'precision_cost' wrt. 'x_sim'.
+
+    :param x_sim:       N x V x T array that contains the simulated time series.
+    :type x_sim:        np.ndarray
+    :param x_target:    N x V x T array that contains the target time series.
+    :type x_target:     np.ndarray
+    :param cost_matrix: N x V binary matrix that defines nodes and channels of precision measurement, defaults to
+                        None
+    :type cost_matrix:  np.ndarray
+    :param interval:    (t_start, t_end). Indices of start and end point of the slice (both inclusive) in time
+                        dimension. Only 'int' positive index-notation allowed (i.e. no negative indices or 'None').
+    :type interval:     tuple
+
+    :return:            Control-dimensions x T array of precision cost gradients.
+    :rtype:             np.ndarray
+    """
+
+    derivative = np.zeros(x_target.shape)
+
+    # integrate over nodes, variables, and time
+    for n in range(x_target.shape[0]):
+        for v in range(x_target.shape[1]):
+            for t in range(interval[0], interval[1]):
+                derivative[n, v, t] = -cost_matrix[n, v] * (x_target[n, v, t] - x_sim[n, v, t])
+
+    return derivative
+
+
+@numba.njit
+def control_strength_cost(u, weights, dt):
+    """Total cost related to the control strength, weighted sum of contributions.
+
+    :param u:           Control-dimensions x T array. Control signals.
+    :type u:            np.ndarray
+    :param weights:     Dictionary of weights.
+    :type weights:      dictionary
+    :param dt:          Time step.
+    :type dt:           float
+
+    :return:            control strength cost of the control.
+    :rtype:             float
+    """
+
+    cost_timeseries = np.zeros((u.shape))
+
+    # timeseries of control vector is weighted sum of contributing cost functionals
+    if weights["w_2"] != 0.0:
+        cost_timeseries += weights["w_2"] * L2_cost(u)
+
+    cost = 0.0
+    # integrate over nodes, channels, and time
+    if weights["w_2"] != 0.0:
+        for n in range(u.shape[0]):
+            for v in range(u.shape[1]):
+                for t in range(u.shape[2]):
+                    cost += cost_timeseries[n, v, t] * dt
+
+    return cost
+
+
+@numba.njit
+def derivative_control_strength_cost(u, weights):
+    """Derivative of the 'control_strength_cost' wrt. the control 'u'.
+
+    :param u:           Control-dimensions x T array. Control signals.
+    :type u:            np.ndarray
+    :param weights:     Dictionary of weights.
+    :type weights:      dictionary
+
+    :return:    Control-dimensions x T array of L2-cost gradients.
+    :rtype:     np.ndarray
+    """
+
+    der = np.zeros((u.shape))
+
+    if weights["w_2"] != 0.0:
+        der += weights["w_2"] * derivative_L2_cost(u)
+
+    return der
+
+
+@numba.njit
+def L2_cost(u):
+    """'Energy' or 'L2' cost. Penalizes for control strength.
+
+    :param u:   Control-dimensions x T array. Control signals.
+    :type u:    np.ndarray
+
+    :return:    L2 cost of the control.
+    :rtype:     float
+    """
+
+    return 0.5 * u**2.0
+
+
+@numba.njit
+def derivative_L2_cost(u):
+    """Derivative of the 'L2_cost' wrt. the control 'u'.
+
+    :param u:   Control-dimensions x T array. Control signals.
+    :type u:    np.ndarray
+
+    :return:    Control-dimensions x T array of L2-cost gradients.
+    :rtype:     np.ndarray
+    """
+    return u