From 4b5e8c6a190ca0c2915b2e264088aa680a8b24d4 Mon Sep 17 00:00:00 2001
From: sbastiaens <sorenza.bastiaens@gmail.com>
Date: Wed, 31 Jul 2024 14:41:22 -0400
Subject: [PATCH 01/15] Folder and files for robinson

---
 whobpyt/models/robinson/robinson_lin.py | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)
 create mode 100644 whobpyt/models/robinson/robinson_lin.py

diff --git a/whobpyt/models/robinson/robinson_lin.py b/whobpyt/models/robinson/robinson_lin.py
new file mode 100644
index 00000000..eca88ba3
--- /dev/null
+++ b/whobpyt/models/robinson/robinson_lin.py
@@ -0,0 +1,20 @@
+"""
+
+WhoBPyt Robinson Linear model classes
+--------------------------------------
+
+Authors: Zheng Wang, John Griffiths, Davide Momi, Sorenza Bastiaens, Parsa Oveisi, Kevin Kadak, Shreyas Harita, Minarose Ismail
+
+Neural Mass Model fitting module for Linear Robinson for M/EEG
+
+"""
+
+
+
+
+"""
+Importage
+---------
+"""
+
+

From 47782ca9ec11203594a7880378529ef7128ba64f Mon Sep 17 00:00:00 2001
From: sbastiaens <62792658+sbastiaens@users.noreply.github.com>
Date: Wed, 31 Jul 2024 14:42:37 -0400
Subject: [PATCH 02/15] Create __init__.py

---
 whobpyt/models/robinson/__init__.py | 1 +
 1 file changed, 1 insertion(+)
 create mode 100644 whobpyt/models/robinson/__init__.py

diff --git a/whobpyt/models/robinson/__init__.py b/whobpyt/models/robinson/__init__.py
new file mode 100644
index 00000000..af7a4d95
--- /dev/null
+++ b/whobpyt/models/robinson/__init__.py
@@ -0,0 +1 @@
+from .robinson import RobinsonModel, RobinsonParams

From d03fc183a5510afd4fc5844a430da5e4415f7a71 Mon Sep 17 00:00:00 2001
From: sbastiaens <62792658+sbastiaens@users.noreply.github.com>
Date: Wed, 31 Jul 2024 15:43:59 -0400
Subject: [PATCH 03/15] Update robinson_lin.py

---
 whobpyt/models/robinson/robinson_lin.py | 173 +++++++++++++++++++++++-
 1 file changed, 171 insertions(+), 2 deletions(-)

diff --git a/whobpyt/models/robinson/robinson_lin.py b/whobpyt/models/robinson/robinson_lin.py
index eca88ba3..3e71773d 100644
--- a/whobpyt/models/robinson/robinson_lin.py
+++ b/whobpyt/models/robinson/robinson_lin.py
@@ -9,12 +9,181 @@
 
 """
 
+"""
+Importage
+---------
+"""
 
+from torch.nn.parameter import Parameter as ptParameter
+from torch.nn import ReLU as ptReLU
+from torch.linalg import norm as ptnorm
+from torch import (tensor as pttensor, float32 as ptfloat32, sum as ptsum, exp as ptexp, diag as ptdiag, 
+                   transpose as pttranspose, zeros_like as ptzeros_like, int64 as ptint64, randn as ptrandn, 
+                   matmul as ptmatmul, tanh as pttanh, matmul as ptmatmul, reshape as ptreshape, sqrt as ptsqrt,
+                   ones as ptones, cat as ptcat)
 
+# Numpy stuff
+from numpy.random import uniform 
+from numpy import ones,zeros
+
+# WhoBPyT stuff
+from ...datatypes import AbstractNeuralModel, AbstractParams, Parameter as par
+from ...functions.arg_type_check import method_arg_type_check
 
 """
-Importage
----------
+Robinson Linear params class
+---------------
+"""
+
+class RobinsonLinParams(AbstractParams):
+    """
+    A class for setting the parameters of a neural mass model for M/EEG data fitting.
+
+    Attributes:
+      Gee (par): Gain of excitatory-excitatory
+      Gei (par): Gain of excitatory-inhibitory
+      Gese (par): Gain of excitatory-relay nuclei-excitatory
+      Gesre (par): Gain of excitatory-relay nuclei-reticular nucleus-excitatory
+      Gsrs (par): Gain of relay nuclei-reticular nucleus-relay nuclei
+      alpha (par): Inverse decay time
+      beta (par): Inverse rise time 
+      t0: Corticothalamic delay
+      EMG: Artifacts
+
+      Not fitted params
+      k0 (par): 10;  Volume conduction parameter 
+      phin (par): 1e-5; Input
+      kmax (par) 4;
+    """
+    def __init__(self, **kwargs):
+        """
+        Initializes the RobinsonLinParams object.
+
+        Args:
+            **kwargs: Keyword arguments for the model parameters.
+
+        Returns:
+            None
+        """
+        # Initializing EC Abeysuriya 2015
+        param = {
+            "Gee": par(2.07),
+            "Gei": par(-4.11),            
+            "Gese": par(0.77*7.77),
+            "Gesre": par(-3.30*0.77*0.66),
+            "Gsrs": par(-3.30*0.20),
+            "alpha": par(83),
+            "beta": par(769),
+            "t0": par(0.085),
+            "EMG": par(0),   
+            "k0": par(10),
+            "phin": par(1e-5)
+            "kmax": par(4)
+         }
+        
+         for var in param:
+            setattr(self, var, param[var])
+
+         for var in kwargs:
+            setattr(self, var, kwargs[var])
+
 """
+RobinsonLin model class
+--------------
+"""
+
+class RobinsonLinModel(AbstractNeuralModel):
+    """
+    A module for forward model (Robinson) to simulate M/EEG power spectra
+    
+    Attibutes
+    ---------
+    state_size : int
+        Equal to 1 for phie: Number of states in the RobinsonLin model
+        
+    output_size : int
+        Equal to 1: Power spectra of phie
+
+    hidden_size: int
+        Number of step_size for each sampling step
+
+    step_size: float
+        Integration step for forward model
+
+    use_fit_gains: bool
+        Flag for fitting gains. 1: fit, 0: not fit
+
+    params: RobinsonLinParams
+        Model parameters object.
+
+    Methods
+    -------
+    createIC(self, ver):
+        Creates the initial conditions for the model.
+
+    setModelParameters(self):    
+        Sets the parameters of the model.
+    
+    forward(params)
+        Forward pass for generating power spectra with current model parameters
+    """
+
+    def __init__(self, 
+                 params: RobinsonLinParams, 
+                 step_size=0.0001, 
+                 output_size=1, 
+                 use_fit_gains=True
+                 ):               
+        """
+        Parameters
+        ----------
+        step_size: float
+            Integration step for forward model
+        output_size : int
+            Output of phie
+        use_fit_gains: bool
+            Flag for fitting gains. 1: fit, 0: not fit
+        params: RobinsonLinParams
+            Model parameters object.
+        """
+        method_arg_type_check(self.__init__) # Check that the passed arguments (excluding self) abide by their expected data types
+        
+        super(RobinsonLinModel, self).__init__(params)
+        self.state_names = ['phie']
+        self.output_names = ["power_spectra"]
+        self.track_params = [] #Is populated during setModelParameters()
+        
+        self.model_name = "RobinsonLin"
+        self.state_size = 1 # Power spectra phie
+        self.step_size = pttensor(step_size, dtype=ptfloat32)  # integration step 0.1 ms
+        self.output_size = output_size  # number of power spectra
+        self.use_fit_gains = use_fit_gains  # flag for fitting gainsfm
+        self.params = params
+        
+        self.setModelParameters()
+        self.setModelSCParameters()  
+
+  
+    def createIC(self, ver, state_lb = -0.5, state_ub = 0.5):
+        """
+        Creates the initial conditions for the model.
+
+        Parameters
+        ----------
+        ver : int # TODO: ADD MORE DESCRIPTION
+            Initial condition version. (in the Robinson model, the version is not used. It is just for consistency with other models)
 
+        Returns
+        -------
+        torch.Tensor
+            Tensor of shape (node_size, state_size) with random values between `state_lb` and `state_ub`.
+        """
+        n_states = self.state_size
+        init_conds = uniform(state_lb, state_ub, (n_nodes, n_states))
+        ptinit_conds = pttensor(init_conds, dtype=ptfloat32)
+                             
+        return ptinit_conds
 
+  
+    def forward(self, external, hx, hE):   
+      # STILL TO DO 

From 50e0cf1f8ceb8fef788c8401c0bd14b2c5b19d58 Mon Sep 17 00:00:00 2001
From: sbastiaens <62792658+sbastiaens@users.noreply.github.com>
Date: Wed, 31 Jul 2024 15:45:53 -0400
Subject: [PATCH 04/15] Create README.md

---
 whobpyt/models/robinson/README.md | 6 ++++++
 1 file changed, 6 insertions(+)
 create mode 100644 whobpyt/models/robinson/README.md

diff --git a/whobpyt/models/robinson/README.md b/whobpyt/models/robinson/README.md
new file mode 100644
index 00000000..7150b17f
--- /dev/null
+++ b/whobpyt/models/robinson/README.md
@@ -0,0 +1,6 @@
+NEEDS TO BE DEVELOPED
+# Robinson Model 
+
+## Robinson Linear model
+
+### Equations

From 5b462bc6468aa197a6406540dfa3fedd8aa21bf9 Mon Sep 17 00:00:00 2001
From: sbastiaens <62792658+sbastiaens@users.noreply.github.com>
Date: Fri, 2 Aug 2024 11:58:33 +0200
Subject: [PATCH 05/15] Update robinson_lin.py

---
 whobpyt/models/robinson/robinson_lin.py | 80 +++++++++++++++++++++++--
 1 file changed, 76 insertions(+), 4 deletions(-)

diff --git a/whobpyt/models/robinson/robinson_lin.py b/whobpyt/models/robinson/robinson_lin.py
index 3e71773d..5742fc76 100644
--- a/whobpyt/models/robinson/robinson_lin.py
+++ b/whobpyt/models/robinson/robinson_lin.py
@@ -42,6 +42,12 @@ class RobinsonLinParams(AbstractParams):
     Attributes:
       Gee (par): Gain of excitatory-excitatory
       Gei (par): Gain of excitatory-inhibitory
+      Ges (par): Gain of excitatory-relay nuclei
+      Gse (par): Gain of relay nuclei-excitatory
+      Gsr (par): Gain of relay nuclei-reticular nucleus
+      Grs (par): Gain of reticular nucleus-relay nulei
+      Gre (par): Gain of reticular nucleus-excitatory
+      Gsn (par): Gain of relay nuclei-input
       Gese (par): Gain of excitatory-relay nuclei-excitatory
       Gesre (par): Gain of excitatory-relay nuclei-reticular nucleus-excitatory
       Gsrs (par): Gain of relay nuclei-reticular nucleus-relay nuclei
@@ -51,9 +57,13 @@ class RobinsonLinParams(AbstractParams):
       EMG: Artifacts
 
       Not fitted params
+      gammae: 116
       k0 (par): 10;  Volume conduction parameter 
+      re (par):
       phin (par): 1e-5; Input
       kmax (par) 4;
+      re (par): 
+      Lx (par): 
     """
     def __init__(self, **kwargs):
         """
@@ -68,17 +78,27 @@ def __init__(self, **kwargs):
         # Initializing EC Abeysuriya 2015
         param = {
             "Gee": par(2.07),
-            "Gei": par(-4.11),            
+            "Gei": par(-4.11),
+            "Ges": par(0.77),
+            "Gse": par(7.77),
+            "Gsr": par(-3.30),
+            "Grs": par(0.20),
+            "Gre": par(0.66),
+            "Gsn": par(8.10),
             "Gese": par(0.77*7.77),
             "Gesre": par(-3.30*0.77*0.66),
             "Gsrs": par(-3.30*0.20),
+            "G_esn": par(0.77*8.10), 
             "alpha": par(83),
             "beta": par(769),
+            "gammae": par(116),
             "t0": par(0.085),
+            "re": par(0.086),
             "EMG": par(0),   
             "k0": par(10),
-            "phin": par(1e-5)
-            "kmax": par(4)
+            "phin": par(1e-5),
+            "kmax": par(4),
+            "Lx": par(0.5)
          }
         
          for var in param:
@@ -185,5 +205,57 @@ def createIC(self, ver, state_lb = -0.5, state_ub = 0.5):
         return ptinit_conds
 
   
-    def forward(self, external, hx, hE):   
+    def forward(self, external, w1):   
       # STILL TO DO 
+      G_ei = self.params.Gei.value()
+      G_ee = self.params.Gee.value()
+      G_es = self.params.Ges.value()
+      G_sn = self.params.Gsn.value()
+      G_se = self.params.Gse.value()
+      G_sr = self.params.Gsr.value()
+      G_rs = self.params.Grs.value()
+      G_re = self.params.Gre.value()
+      alpha = 83
+      beta = 769
+      G_esn = G_es*G_sn
+      G_srs = G_sr*G_rs
+      G_esre = G_es*G_sr*G_re
+      G_ese = G_es*G_se
+      gamma_e = self.params.gammae.value()
+      t0 = self.params.t0.value()
+      r_e = self.params.re.value()
+      phin = self.params.phin.value()
+      EMG = self.params.EMG.value()
+      kmax = self.params.kmax.value()
+      k0 = self.params.k0.value()
+      Lx = self.params.Lx.value()
+
+      dk = 2*np.pi/Lx
+      m_rows = list(range(-kmax, kmax + 1))
+      n_cols = list(range(-kmax, kmax + 1))
+      [kxa,kya] = np.meshgrid(dk*np.array(m_rows),dk*np.array(n_cols))
+      k2 = kxa**(2)+kya**(2)
+      k2u = np.unique(k2[:])
+      counts, _ = np.histogram(k2, bins=np.append(k2u, k2u[-1] + 1))
+      k2u = np.column_stack((k2u, counts))
+      k2_volconduct = np.exp(-k2u[:,0]/k0**2);
+      emg_f = 40;
+      emg = ((w1/(2*np.pi*emg_f))**2)/((1+(w1/(2*np.pi*emg_f))**2)**2);
+      if EMG==0:
+        emg = 0;
+    
+      L = (1-(w1*1j/alpha))**(-1)*(1-(w1*1j/beta))**(-1)
+      phin = 1e-5
+      Gei_oneminus = 1-L*G_ei
+      Gsrs_oneminus = 1-G_srs*(L**(2))
+      re2 = r_e**2
+      q2re2 = (1-(w1*1j/gamma_e))**(2) - (1/(1-G_ei*L))*(G_ee*L + ((G_ese*L**(2) + G_esre*L**(3))*np.exp(w1*1j*t0))/(1-G_srs*L**(2)))
+      T_prefactor = (L**(2)*phin)/(Gei_oneminus*Gsrs_oneminus);
+    
+      P = np.zeros_like(w1)
+      for j in range(k2u.shape[0]):
+          contribution = k2u[j, 1] * np.abs(T_prefactor / (k2u[j, 0] * re2 + q2re2))**2 * k2_volconduct[j]
+          P += contribution
+    
+      P = P + 1e-12*0.1*emg;
+      return P

From 8cec7bfb6fae2a9fc5783bfebc6b4edb7038f10c Mon Sep 17 00:00:00 2001
From: sbastiaens <62792658+sbastiaens@users.noreply.github.com>
Date: Wed, 7 Aug 2024 08:16:34 -0400
Subject: [PATCH 06/15] Update robinson_lin.py

---
 whobpyt/models/robinson/robinson_lin.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/whobpyt/models/robinson/robinson_lin.py b/whobpyt/models/robinson/robinson_lin.py
index 5742fc76..9100b8ad 100644
--- a/whobpyt/models/robinson/robinson_lin.py
+++ b/whobpyt/models/robinson/robinson_lin.py
@@ -206,7 +206,7 @@ def createIC(self, ver, state_lb = -0.5, state_ub = 0.5):
 
   
     def forward(self, external, w1):   
-      # STILL TO DO 
+      # generate power spectra
       G_ei = self.params.Gei.value()
       G_ee = self.params.Gee.value()
       G_es = self.params.Ges.value()

From c48d6411083e160157eec2c15a0b267d7b1b9387 Mon Sep 17 00:00:00 2001
From: sbastiaens <62792658+sbastiaens@users.noreply.github.com>
Date: Wed, 7 Aug 2024 08:22:40 -0400
Subject: [PATCH 07/15] Update __init__.py

---
 whobpyt/models/robinson/__init__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/whobpyt/models/robinson/__init__.py b/whobpyt/models/robinson/__init__.py
index af7a4d95..81dfd2c6 100644
--- a/whobpyt/models/robinson/__init__.py
+++ b/whobpyt/models/robinson/__init__.py
@@ -1 +1 @@
-from .robinson import RobinsonModel, RobinsonParams
+from .robinson import RobinsonLinModel, RobinsonLinParams

From 1e737cdfeb581c322d6ab3a57a2309d12e869172 Mon Sep 17 00:00:00 2001
From: sbastiaens <62792658+sbastiaens@users.noreply.github.com>
Date: Wed, 7 Aug 2024 14:04:28 +0000
Subject: [PATCH 08/15] Forward function Robinson Linear example

---
 examples/robinson_lin_forward.py | 195 +++++++++++++++++++++++++++++++
 1 file changed, 195 insertions(+)
 create mode 100644 examples/robinson_lin_forward.py

diff --git a/examples/robinson_lin_forward.py b/examples/robinson_lin_forward.py
new file mode 100644
index 00000000..4835b950
--- /dev/null
+++ b/examples/robinson_lin_forward.py
@@ -0,0 +1,195 @@
+# Reproducing Abeysuriya 2015 outputs and BrainTrak equations with the forward function
+
+# Importage:
+
+# os stuff
+import os
+import sys
+sys.path.append('..')
+
+# whobpyt stuff
+import whobpyt
+from whobpyt.datatypes import Parameter as par
+from whobpyt.models.robinson import RobinsonLinModel, RobinsonLinParams
+
+# python stuff
+import numpy as np
+import pandas as pd
+import scipy.io
+import pickle
+import warnings
+warnings.filterwarnings('ignore')
+import math
+import numpy as np
+#neuroimaging packages
+import mne
+
+# viz stuff
+import matplotlib.pyplot as plt
+
+def parameter_set(pars):
+  if pars == 'EC':
+    # Eyes closed
+    params = RobinsonLinParams( Gei = par(-4.11),
+                                Gee = par(2.070),
+                                Ges = par(0.77),
+                                Gsn = par(8.10),
+                                Gse = par(7.77),
+                                Gsr = par(-3.3),
+                                Grs = par(0.2),
+                                Gre = par(0.66),
+                                alpha = par(83),
+                                beta = par(769),
+                                gammae = par(116),
+                                t0 = par(0.085),
+                                re = par(0.086),
+                                phin = par(1e-5),
+                                EMG = par(0),
+                                kmax = par(4),
+                                k0 = par(10),
+                                Lx = par(0.5)
+                                )
+
+  if pars == 'EO':
+    # Eyes open
+    params = RobinsonLinParams( Gei = par(-13.22),
+                                Gee = par(10.50),
+                                Ges = par(1.21),
+                                Gsn = par(14.23),
+                                Gse = par(5.78),
+                                Gsr = par(-2.83),
+                                Grs = par(0.25),
+                                Gre = par(0.85),
+                                alpha = par(83),
+                                beta = par(769),
+                                gammae = par(116),
+                                t0 = par(0.085),
+                                re = par(0.086),
+                                phin = par(1e-5),
+                                EMG = par(0),
+                                kmax = par(4),
+                                k0 = par(10),
+                                Lx = par(0.5)
+                                )      
+  if pars == 'REM':
+    # REM
+    params = RobinsonLinParams( Gei = par(-6.61),
+                                Gee = par(5.87),
+                                Ges = par(0.21),
+                                Gsn = par(0.68),
+                                Gse = par(0.66),
+                                Gsr = par(-0.28),
+                                Grs = par(4.59),
+                                Gre = par(2.08),
+                                alpha = par(45),
+                                beta = par(185),
+                                gammae = par(116),
+                                t0 = par(0.085),
+                                re = par(0.086),
+                                phin = par(1e-5),
+                                EMG = par(0),
+                                kmax = par(4),
+                                k0 = par(10),
+                                Lx = par(0.5)
+                                )   
+  if pars == 'S1':
+    # S1
+    params = RobinsonLinParams( Gei = par(-8.30),
+                                Gee = par(7.45),
+                                Ges = par(0.31),
+                                Gsn = par(3.90),
+                                Gse = par(1.67),
+                                Gsr = par(-0.40),
+                                Grs = par(4.44),
+                                Gre = par(7.47),
+                                alpha = par(45),
+                                beta = par(185),
+                                gammae = par(116),
+                                t0 = par(0.085),
+                                re = par(0.086),
+                                phin = par(1e-5),
+                                EMG = par(0),
+                                kmax = par(4),
+                                k0 = par(10),
+                                Lx = par(0.5)
+                                ) 
+  if pars == 'S2':
+    # S2
+    params = RobinsonLinParams( Gei = par(-17.93),
+                                Gee = par(16.8),
+                                Ges = par(3.89),
+                                Gsn = par(2.38),
+                                Gse = par(0.07),
+                                Gsr = par(-0.14),
+                                Grs = par(8.33),
+                                Gre = par(4.96),
+                                alpha = par(45),
+                                beta = par(185),
+                                gammae = par(116),
+                                t0 = par(0.085),
+                                re = par(0.086),
+                                phin = par(1e-5),
+                                EMG = par(0),
+                                kmax = par(4),
+                                k0 = par(10),
+                                Lx = par(0.5)
+                                )  
+  if pars == 'SWS':
+    # SWS
+    params = RobinsonLinParams( Gei = par(-19.74),
+                                Gee = par(19.52),
+                                Ges = par(5.30),
+                                Gsn = par(1.70),
+                                Gse = par(0.22),
+                                Gsr = par(-0.22),
+                                Grs = par(1.35),
+                                Gre = par(1.90),
+                                alpha = par(45),
+                                beta = par(185),
+                                gammae = par(116),
+                                t0 = par(0.085),
+                                re = par(0.086),
+                                phin = par(1e-5),
+                                EMG = par(0),
+                                kmax = par(4),
+                                k0 = par(10),
+                                Lx = par(0.5)
+                                )   
+  if pars == 'Spindles':
+    # Spindles
+    params = RobinsonLinParams( Gei = par(-18.96),
+                                Gee = par(18.52),
+                                Ges = par(2.55),
+                                Gsn = par(2.78),
+                                Gse = par(0.73),
+                                Gsr = par(-0.26),
+                                Grs = par(16.92),
+                                Gre = par(4.67),
+                                alpha = par(45),
+                                beta = par(185),
+                                gammae = par(116),
+                                t0 = par(0.085),
+                                re = par(0.086),
+                                phin = par(1e-5),
+                                EMG = par(0),
+                                kmax = par(4),
+                                k0 = par(10),
+                                Lx = par(0.5)
+                                )   
+  return params
+
+freq = np.linspace(1,60,200)
+w1 = 2*math.pi*freq
+options = ['EC','EO', 'REM', 'S1', 'S2', 'SWS', 'Spindles']
+for option in options:
+    params = parameter_set(option)
+    model = RobinsonLinModel(params)
+    P = model.forward(external=0, w1=w1)
+    plt.plot(freq,P/sum(P))
+plt.yscale("log")
+plt.xscale("log")
+plt.ylabel('Normalized Power (A.U.)')
+plt.xlabel('Frequency (Hz)')
+plt.legend(options)
+plt.savefig('plot.png')
+plt.title('Power spectra for different parameter sets')
\ No newline at end of file

From 19b1387ce8399cdc4fe89130147d60883b309847 Mon Sep 17 00:00:00 2001
From: sbastiaens <62792658+sbastiaens@users.noreply.github.com>
Date: Wed, 7 Aug 2024 10:09:28 -0400
Subject: [PATCH 09/15] Update robinson_lin_forward.py

---
 examples/robinson_lin_forward.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/examples/robinson_lin_forward.py b/examples/robinson_lin_forward.py
index 4835b950..de060935 100644
--- a/examples/robinson_lin_forward.py
+++ b/examples/robinson_lin_forward.py
@@ -191,5 +191,4 @@ def parameter_set(pars):
 plt.ylabel('Normalized Power (A.U.)')
 plt.xlabel('Frequency (Hz)')
 plt.legend(options)
-plt.savefig('plot.png')
-plt.title('Power spectra for different parameter sets')
\ No newline at end of file
+plt.title('Power spectra for different parameter sets')

From bbfa048d7f7a3d341c186b95b15e24c6697ce204 Mon Sep 17 00:00:00 2001
From: sbastiaens <62792658+sbastiaens@users.noreply.github.com>
Date: Wed, 7 Aug 2024 10:29:27 -0400
Subject: [PATCH 10/15] Sphinx gallery line

---
 examples/robinson_lin_forward.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/examples/robinson_lin_forward.py b/examples/robinson_lin_forward.py
index de060935..53bd7449 100644
--- a/examples/robinson_lin_forward.py
+++ b/examples/robinson_lin_forward.py
@@ -1,5 +1,6 @@
 # Reproducing Abeysuriya 2015 outputs and BrainTrak equations with the forward function
 
+# sphinx_gallery_thumbnail_number = 2
 # Importage:
 
 # os stuff

From 4dd2cbf357b9684dcad6bdd142865c1e73e780cd Mon Sep 17 00:00:00 2001
From: sbastiaens <62792658+sbastiaens@users.noreply.github.com>
Date: Wed, 7 Aug 2024 10:39:35 -0400
Subject: [PATCH 11/15] Update robinson_lin_forward.py

---
 examples/robinson_lin_forward.py | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/examples/robinson_lin_forward.py b/examples/robinson_lin_forward.py
index 53bd7449..0542a023 100644
--- a/examples/robinson_lin_forward.py
+++ b/examples/robinson_lin_forward.py
@@ -1,3 +1,20 @@
+"""
+.. _ex-robinson_lin:
+
+========================================================
+Testing forward function of Linear Robinson model
+========================================================
+
+This example runs a minimal example for running the forward function of the Linear Robinson based on 
+
+R.G. Abeysuriya, C.J. Rennie, P.A. Robinson (2015)
+"Physiologically based arousal state estimation and dynamics"
+
+
+The code includes data fetching, model fitting, and result visualization based on the methods presented in the paper.
+
+"""
+
 # Reproducing Abeysuriya 2015 outputs and BrainTrak equations with the forward function
 
 # sphinx_gallery_thumbnail_number = 2

From ea5bb5736bfe1272663d58e32a89d6eff020059a Mon Sep 17 00:00:00 2001
From: sbastiaens <62792658+sbastiaens@users.noreply.github.com>
Date: Thu, 8 Aug 2024 15:46:28 +0000
Subject: [PATCH 12/15] Updated version

---
 whobpyt/models/robinson/__init__.py     |  4 ++++
 whobpyt/models/robinson/robinson_lin.py | 30 +++++++++++++++----------
 2 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/whobpyt/models/robinson/__init__.py b/whobpyt/models/robinson/__init__.py
index 81dfd2c6..397382c0 100644
--- a/whobpyt/models/robinson/__init__.py
+++ b/whobpyt/models/robinson/__init__.py
@@ -1 +1,5 @@
+<<<<<<< Updated upstream
 from .robinson import RobinsonLinModel, RobinsonLinParams
+=======
+from .robinson_lin import RobinsonLinModel, RobinsonLinParams
+>>>>>>> Stashed changes
diff --git a/whobpyt/models/robinson/robinson_lin.py b/whobpyt/models/robinson/robinson_lin.py
index 9100b8ad..3a0c1dee 100644
--- a/whobpyt/models/robinson/robinson_lin.py
+++ b/whobpyt/models/robinson/robinson_lin.py
@@ -20,11 +20,12 @@
 from torch import (tensor as pttensor, float32 as ptfloat32, sum as ptsum, exp as ptexp, diag as ptdiag, 
                    transpose as pttranspose, zeros_like as ptzeros_like, int64 as ptint64, randn as ptrandn, 
                    matmul as ptmatmul, tanh as pttanh, matmul as ptmatmul, reshape as ptreshape, sqrt as ptsqrt,
-                   ones as ptones, cat as ptcat)
+                   ones as ptones, cat as ptcat, pi as ptpi, abs as ptabs)
 
 # Numpy stuff
 from numpy.random import uniform 
 from numpy import ones,zeros
+import numpy as np
 
 # WhoBPyT stuff
 from ...datatypes import AbstractNeuralModel, AbstractParams, Parameter as par
@@ -101,10 +102,10 @@ def __init__(self, **kwargs):
             "Lx": par(0.5)
          }
         
-         for var in param:
+        for var in param:
             setattr(self, var, param[var])
 
-         for var in kwargs:
+        for var in kwargs:
             setattr(self, var, kwargs[var])
 
 """
@@ -181,7 +182,6 @@ def __init__(self,
         self.params = params
         
         self.setModelParameters()
-        self.setModelSCParameters()  
 
   
     def createIC(self, ver, state_lb = -0.5, state_ub = 0.5):
@@ -215,8 +215,8 @@ def forward(self, external, w1):
       G_sr = self.params.Gsr.value()
       G_rs = self.params.Grs.value()
       G_re = self.params.Gre.value()
-      alpha = 83
-      beta = 769
+      alpha = self.params.alpha.value()
+      beta = self.params.beta.value()
       G_esn = G_es*G_sn
       G_srs = G_sr*G_rs
       G_esre = G_es*G_sr*G_re
@@ -230,17 +230,23 @@ def forward(self, external, w1):
       k0 = self.params.k0.value()
       Lx = self.params.Lx.value()
 
-      dk = 2*np.pi/Lx
+      w1 = pttensor(w1)  
+      dk = 2*ptpi/Lx
+      kmax = int(kmax.item())
       m_rows = list(range(-kmax, kmax + 1))
       n_cols = list(range(-kmax, kmax + 1))
       [kxa,kya] = np.meshgrid(dk*np.array(m_rows),dk*np.array(n_cols))
+      kxa = pttensor(kxa)
+      kya = pttensor(kya)
+
       k2 = kxa**(2)+kya**(2)
       k2u = np.unique(k2[:])
       counts, _ = np.histogram(k2, bins=np.append(k2u, k2u[-1] + 1))
-      k2u = np.column_stack((k2u, counts))
-      k2_volconduct = np.exp(-k2u[:,0]/k0**2);
+      k2u = pttensor(np.column_stack((k2u, counts)))
+      k2_volconduct = pttensor(np.exp(-k2u[:,0]/k0**2));
       emg_f = 40;
       emg = ((w1/(2*np.pi*emg_f))**2)/((1+(w1/(2*np.pi*emg_f))**2)**2);
+      emg = pttensor(emg)
       if EMG==0:
         emg = 0;
     
@@ -249,12 +255,12 @@ def forward(self, external, w1):
       Gei_oneminus = 1-L*G_ei
       Gsrs_oneminus = 1-G_srs*(L**(2))
       re2 = r_e**2
-      q2re2 = (1-(w1*1j/gamma_e))**(2) - (1/(1-G_ei*L))*(G_ee*L + ((G_ese*L**(2) + G_esre*L**(3))*np.exp(w1*1j*t0))/(1-G_srs*L**(2)))
+      q2re2 = (1-(w1*1j/gamma_e))**(2) - (1/(1-G_ei*L))*(G_ee*L + ((G_ese*L**(2) + G_esre*L**(3))*ptexp(w1*1j*t0))/(1-G_srs*L**(2)))
       T_prefactor = (L**(2)*phin)/(Gei_oneminus*Gsrs_oneminus);
     
-      P = np.zeros_like(w1)
+      P = pttensor(np.zeros_like(w1))
       for j in range(k2u.shape[0]):
-          contribution = k2u[j, 1] * np.abs(T_prefactor / (k2u[j, 0] * re2 + q2re2))**2 * k2_volconduct[j]
+          contribution = k2u[j, 1] * ptabs(T_prefactor / (k2u[j, 0] * re2 + q2re2))**2 * k2_volconduct[j]
           P += contribution
     
       P = P + 1e-12*0.1*emg;

From faa883bdf357b909b4ed081a4ddd78c990b3d9f5 Mon Sep 17 00:00:00 2001
From: sbastiaens <62792658+sbastiaens@users.noreply.github.com>
Date: Wed, 4 Dec 2024 14:54:17 -0500
Subject: [PATCH 13/15] Update __init__.py

---
 whobpyt/models/robinson/__init__.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/whobpyt/models/robinson/__init__.py b/whobpyt/models/robinson/__init__.py
index 397382c0..d9017577 100644
--- a/whobpyt/models/robinson/__init__.py
+++ b/whobpyt/models/robinson/__init__.py
@@ -1,5 +1 @@
-<<<<<<< Updated upstream
-from .robinson import RobinsonLinModel, RobinsonLinParams
-=======
 from .robinson_lin import RobinsonLinModel, RobinsonLinParams
->>>>>>> Stashed changes

From e84b2bf12ac42b3b772a61aedbd438a532c70379 Mon Sep 17 00:00:00 2001
From: sbastiaens <62792658+sbastiaens@users.noreply.github.com>
Date: Wed, 4 Dec 2024 16:14:01 -0500
Subject: [PATCH 14/15] Update index.rst

---
 doc/index.rst | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/doc/index.rst b/doc/index.rst
index a500acba..1bc54e7f 100644
--- a/doc/index.rst
+++ b/doc/index.rst
@@ -33,9 +33,10 @@ Welcome to WhoBPyT documentation!
    API/datatypes
 
 .. toctree::
-   :maxdepth: 1
+   :maxdepth: 2
    :caption: Examples
    :glob:
  
    auto_examples/eg__tmseeg
+   auto_examples/eg__robinson_lin_forward
    

From 1e851df40a453d0c1c6e1c59276f9f065085a8a0 Mon Sep 17 00:00:00 2001
From: sbastiaens <62792658+sbastiaens@users.noreply.github.com>
Date: Wed, 4 Dec 2024 16:20:35 -0500
Subject: [PATCH 15/15] Update and rename robinson_lin_forward.py to
 eg__robinson_lin_forward.py

---
 .../{robinson_lin_forward.py => eg__robinson_lin_forward.py}   | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)
 rename examples/{robinson_lin_forward.py => eg__robinson_lin_forward.py} (99%)

diff --git a/examples/robinson_lin_forward.py b/examples/eg__robinson_lin_forward.py
similarity index 99%
rename from examples/robinson_lin_forward.py
rename to examples/eg__robinson_lin_forward.py
index 0542a023..7f09e659 100644
--- a/examples/robinson_lin_forward.py
+++ b/examples/eg__robinson_lin_forward.py
@@ -195,7 +195,8 @@ def parameter_set(pars):
                                 Lx = par(0.5)
                                 )   
   return params
-
+# %%
+# Plot
 freq = np.linspace(1,60,200)
 w1 = 2*math.pi*freq
 options = ['EC','EO', 'REM', 'S1', 'S2', 'SWS', 'Spindles']