Code Standardization

doc/API/models.rst

@@ -11,7 +11,7 @@ Models
    :undoc-members:
    :special-members: __init__
 
-.. autoclass:: whobpyt.models.RWW2.RWW2
+.. autoclass:: whobpyt.models.RWWEI2.RWWEI2
    :members: info, createIC, setModelParameters, forward
    :undoc-members:
    :special-members: __init__

examples/eg002r__multimodal_simulation.py

@@ -24,7 +24,7 @@
 # whobpyt stuff
 import whobpyt
 from whobpyt.datatypes import par, Recording
-from whobpyt.models.RWW2 import mmRWW2, mmRWW2_np, RWW2, RWW2_np, ParamsRWW2
+from whobpyt.models.RWWEI2 import RWWEI2_EEG_BOLD, RWWEI2_EEG_BOLD_np, RWWEI2, RWWEI2_np, ParamsRWWEI2
 from whobpyt.models.BOLD import BOLD_Layer, BOLD_np, BOLD_Params
 from whobpyt.models.EEG import EEG_Layer, EEG_np, EEG_Params
 from whobpyt.optimization import CostsFC, CostsPSD, CostsMean, CostsFixedFC, CostsFixedPSD
@@ -67,7 +67,7 @@
 init_state = (init_state + torch.randn_like(init_state)/30).to(device) # Randomizing initial values
 
 # Create a RWW Params
-paramsNode = ParamsRWW2(num_regions)
+paramsNode = ParamsRWWEI2(num_regions)
 
 #Create #EEG Params
 paramsEEG = EEG_Params(torch.eye(num_regions))
@@ -105,7 +105,7 @@
 # The Multi-Modal Model
 
 
-model = mmRWW2(num_regions, num_channels, paramsNode, paramsEEG, paramsBOLD, Con_Mtx, dist_mtx, step_size, sim_len, device = device)
+model = RWWEI2_EEG_BOLD(num_regions, num_channels, paramsNode, paramsEEG, paramsBOLD, Con_Mtx, dist_mtx, step_size, sim_len, device = device)
 
 
 # %%
@@ -135,15 +135,15 @@ def __init__(self):
         #self.BOLD_PSD = CostsPSD(...) # Not Currently Used
         #self.BOLD_FC = CostsFC(num_regions, varIdx = 4, targetValue = SC_mtx_norm)
 
-    def loss(self, node_history, EEG_history, BOLD_history, temp, returnLossComponents = False):
+    def loss(self, simData, empData = None, returnLossComponents = False):
         # sim, ts_window, self.model, next_window
 
-        S_E_mean_loss = self.S_E_mean.calcLoss(node_history) 
-        S_I_mean_loss = torch.tensor([0]).to(device) #self.S_I_mean.calcLoss(node_history)
-        EEG_PSD_loss = torch.tensor([0]).to(device) #self.EEG_PSD.calcLoss(EEG_history) 
-        EEG_FC_loss = torch.tensor([0]).to(device) #self.EEG_FC.calcLoss(EEG_history)
-        BOLD_PSD_loss = torch.tensor([0]).to(device) #self.BOLD_PS.calcLoss(BOLD_history)
-        BOLD_FC_loss = torch.tensor([0]).to(device) #self.BOLD_FC.calcLoss(BOLD_history)
+        S_E_mean_loss = self.S_E_mean.loss(simData) 
+        S_I_mean_loss = torch.tensor([0]).to(device) #self.S_I_mean.loss(simData)
+        EEG_PSD_loss = torch.tensor([0]).to(device) #self.EEG_PSD.loss(simData) 
+        EEG_FC_loss = torch.tensor([0]).to(device) #self.EEG_FC.loss(simData)
+        BOLD_PSD_loss = torch.tensor([0]).to(device) #self.BOLD_PS.loss(simData)
+        BOLD_FC_loss = torch.tensor([0]).to(device) #self.BOLD_FC.loss(simData)
 
         totalLoss = self.S_E_mean_weight*S_E_mean_loss + self.S_I_mean_weight*S_I_mean_loss \
                   + self.EEG_PSD_weight*EEG_PSD_loss   + self.EEG_FC_weight*EEG_FC_loss \
@@ -246,7 +246,7 @@ def loss(self, node_history, EEG_history, BOLD_history, temp, returnLossComponen
 model.eeg.params.LF = model.eeg.params.LF.cpu()
 
 val_sim_len = 20*1000 # Simulation length in msecs
-model_validate = mmRWW2_np(num_regions, num_channels, model.params, model.eeg.params, model.bold.params, Con_Mtx.detach().cpu().numpy(), dist_mtx.detach().cpu().numpy(), step_size, val_sim_len)
+model_validate = RWWEI2_EEG_BOLD_np(num_regions, num_channels, model.params, model.eeg.params, model.bold.params, Con_Mtx.detach().cpu().numpy(), dist_mtx.detach().cpu().numpy(), step_size, val_sim_len)
 
 sim_vals, hE = model_validate.forward(external = 0, hx = model_validate.createIC(ver = 0), hE = 0)
 

examples/eg003r__fitting_rww_example.py

@@ -114,7 +114,7 @@
 
 # %%
 # create objective function
-ObjFun = CostsRWW()
+ObjFun = CostsRWW(model)
 
 # %%
 # call model fit

examples/eg004r__fitting_JR_example.py

@@ -112,7 +112,7 @@
 
 # %%
 # create objective function
-ObjFun = CostsJR()
+ObjFun = CostsJR(model)
 
 # %%
 # call model fit

examples/eg005r__gpu_support.py

@@ -20,11 +20,11 @@
 # whobpyt stuff
 import whobpyt
 from whobpyt.datatypes import par, Recording
-from whobpyt.models.RWW2 import mmRWW2, mmRWW2_np, RWW2, RWW2_np, ParamsRWW2
+from whobpyt.models.RWWEI2 import RWWEI2_EEG_BOLD, RWWEI2_EEG_BOLD_np, RWWEI2, RWWEI2_np, ParamsRWWEI2
 from whobpyt.models.BOLD import BOLD_Layer, BOLD_np, BOLD_Params
 from whobpyt.models.EEG import EEG_Layer, EEG_np, EEG_Params
 from whobpyt.optimization import CostsFC, CostsPSD, CostsMean, CostsFixedFC, CostsFixedPSD
-from whobpyt.optimization.custom_cost_mmRWW2 import CostsmmRWW2
+from whobpyt.optimization.custom_cost_mmRWW2 import CostsmmRWWEI2
 from whobpyt.run import Model_fitting, Fitting_FNGFPG, Fitting_Batch
 from whobpyt.data.generators import gen_cube
 
@@ -67,7 +67,7 @@
 plt.title("SC of Artificial Data")
 
 # Create a RWW Params
-paramsNode = ParamsRWW2(num_regions)
+paramsNode = ParamsRWWEI2(num_regions)
 
 paramsNode.J = par((0.15 * np.ones(num_regions)), fit_par = True, asLog = True) #This is a parameter that will be updated during training
 paramsNode.G = par(torch.tensor(1.0), None, None, True, False, False)
@@ -92,7 +92,7 @@
 # Simulation Length
 step_size = 0.1 # Step Size in msecs
 sim_len = 1500 # Simulation length in msecs
-model = RWW2(num_regions, paramsNode, Con_Mtx, dist_mtx, step_size, sim_len, device = device)
+model = RWWEI2(num_regions, paramsNode, Con_Mtx, dist_mtx, step_size, sim_len, device = device)
 
 demoPSD = torch.rand(100).to(device)
 objFun = CostsFixedPSD(num_regions = num_regions, simKey = "E", sampleFreqHz = 10000, minFreq = 1, maxFreq = 100, targetValue = demoPSD, rmTransient = 5000, device = device)
@@ -126,10 +126,10 @@
 # Simulation Length
 step_size = 0.1 # Step Size in msecs
 sim_len = 5000 # Simulation length in msecs
-model = mmRWW2(num_regions, num_channels, model.params, paramsEEG, paramsBOLD, Con_Mtx, dist_mtx, step_size, sim_len, device)
+model = RWWEI2_EEG_BOLD(num_regions, num_channels, model.params, paramsEEG, paramsBOLD, Con_Mtx, dist_mtx, step_size, sim_len, device)
 
 targetValue = torch.tensor([0.164]).to(device)
-objFun = CostsmmRWW2(num_regions, simKey = "E", targetValue = targetValue, device = device)
+objFun = CostsmmRWWEI2(num_regions, simKey = "E", targetValue = targetValue, device = device)
 
 # Create a Fitting Object
 F = Fitting_FNGFPG(model, objFun, device)
@@ -163,7 +163,7 @@
 model.eeg.params.LF = model.eeg.params.LF.cpu()
 
 val_sim_len = 20*1000 # Simulation length in msecs
-model_validate = mmRWW2_np(num_regions, num_channels, model.params, model.eeg.params, model.bold.params, Con_Mtx.detach().cpu().numpy(), dist_mtx.detach().cpu().numpy(), step_size, val_sim_len)
+model_validate = RWWEI2_EEG_BOLD_np(num_regions, num_channels, model.params, model.eeg.params, model.bold.params, Con_Mtx.detach().cpu().numpy(), dist_mtx.detach().cpu().numpy(), step_size, val_sim_len)
 
 sim_vals, hE = model_validate.forward(external = 0, hx = model_validate.createIC(ver = 0), hE = 0)
 

whobpyt/datatypes/AbstractLoss.py

@@ -3,11 +3,16 @@
 class AbstractLoss:
     # This is the abstract class for objective function components, or for a custom objective function with multiple components. 
 
-    def __init__(self):
+    def __init__(self, simKey = None, device = torch.device('cpu')):
 
-        self.simKey = None #This is a string key to extract from the dictionary of simulation outputs the time series used by the objective function
+        self.simKey = simKey #This is a string key to extract from the dictionary of simulation outputs the time series used by the objective function
+        device = device
 
-    def loss(self, sim, emp, simKey, model: torch.nn.Module, state_vals):
+    def loss(self, simData, empData):
         # Calculates a loss to be backpropagated through
-        # TODO: In some classes this function is called calcLoss, need to make consistent
+        # If the objective function needs additional info, it should be defined at initialization so that the parameter fitting paradigms don't need to change
+
+        # simData: is a dictionary of simulated state variable/neuroimaging modality time series. Typically accessed as simData[self.simKey].
+        # empData: is the target either as a time series or a calculated phenomena metric
+
         pass

whobpyt/models/RWW2/mmRWW2.py → whobpyt/models/RWWEI2/Multimodal_RWWEI2.py

@@ -2,19 +2,19 @@
 
 import torch
 from whobpyt.datatypes import AbstractNMM, AbstractMode, AbstractParams
-from whobpyt.models.RWW2 import RWW2, ParamsRWW2
+from whobpyt.models.RWWEI2 import RWWEI2, ParamsRWWEI2
 from whobpyt.models.BOLD import BOLD_Layer, BOLD_Params
 from whobpyt.models.EEG import EEG_Layer, EEG_Params
 
-class mmRWW2(RWW2):
+class RWWEI2_EEG_BOLD(RWWEI2):
 
-    model_name = "mmRWW2"
+    model_name = "RWWEI2_EEG_BOLD"
 
     def __init__(self, num_regions, num_channels, paramsNode, paramsEEG, paramsBOLD, Con_Mtx, dist_mtx, step_size, sim_len, device = torch.device('cpu')):
 
         self.eeg = EEG_Layer(num_regions, paramsEEG, num_channels, device = device)
         self.bold = BOLD_Layer(num_regions, paramsBOLD, device = device)
-        super(mmRWW2, self).__init__(num_regions, paramsNode, Con_Mtx, dist_mtx, step_size, useBC = False, device = device)
+        super(RWWEI2_EEG_BOLD, self).__init__(num_regions, paramsNode, Con_Mtx, dist_mtx, step_size, useBC = False, device = device)
 
         self.node_size = num_regions
         self.step_size = step_size
@@ -52,12 +52,12 @@ def forward(self, external, hx, hE, setNoise = None):
         return forward(self, external, hx, hE, setNoise)
 
 def setModelParameters(self):
-    super(mmRWW2, self).setModelParameters() # Currently this is the only one with parameters being fitted
+    super(RWWEI2_EEG_BOLD, self).setModelParameters() # Currently this is the only one with parameters being fitted
     self.eeg.setModelParameters()
     self.bold.setModelParameters()    
 
 def createIC(self, ver):
-    #super(mmRWW2, self).createIC()
+    #super(RWWEI2_EEG_BOLD, self).createIC()
     #self.eeg.createIC()
     #self.bold.createIC()
 
@@ -69,7 +69,7 @@ def createIC(self, ver):
 
 def forward(self, external, hx, hE, setNoise):
 
-    NMM_vals, hE = super(mmRWW2, self).forward(external, self.next_start_state[:, 0:2, :], hE, setNoise) #TODO: Fix the hx in the future
+    NMM_vals, hE = super(RWWEI2_EEG_BOLD, self).forward(external, self.next_start_state[:, 0:2, :], hE, setNoise) #TODO: Fix the hx in the future
     EEG_vals, hE = self.eeg.forward(self.step_size, self.sim_len, NMM_vals["E"].permute((1,0,2)))
     BOLD_vals, hE = self.bold.forward(self.next_start_state[:, 2:6, :], self.step_size, self.sim_len, NMM_vals["E"].permute((1,0,2)))
 

whobpyt/models/RWW2/mmRWW2_validate.py → whobpyt/models/RWWEI2/Multimodal_RWWEI2_validate.py

@@ -2,16 +2,16 @@
 
 import numpy as np
 from whobpyt.datatypes import AbstractNMM, AbstractMode, AbstractParams
-from whobpyt.models.RWW2 import RWW2_np, ParamsRWW2
+from whobpyt.models.RWWEI2 import RWWEI2_np, ParamsRWWEI2
 from whobpyt.models.BOLD import BOLD_np, BOLD_Params
 from whobpyt.models.EEG import EEG_np, EEG_Params
 
-class mmRWW2_np(RWW2_np):
+class RWWEI2_EEG_BOLD_np(RWWEI2_np):
 
-    model_name = "mmRWW2_np"
+    model_name = "RWWEI2_EEG_BOLD_np"
 
     def __init__(self, num_regions, num_channels, paramsNode, paramsEEG, paramsBOLD, Con_Mtx, dist_mtx, step_size, sim_len):
-        super(mmRWW2_np, self).__init__(num_regions, paramsNode, Con_Mtx, dist_mtx, step_size)
+        super(RWWEI2_EEG_BOLD_np, self).__init__(num_regions, paramsNode, Con_Mtx, dist_mtx, step_size)
         self.eeg = EEG_np(num_regions, paramsEEG, num_channels)
         self.bold = BOLD_np(num_regions, paramsBOLD)
 
@@ -39,12 +39,12 @@ def forward(self, external, hx, hE):
         return forward(self, external, hx, hE)
 
 def setModelParameters(self):
-    super(mmRWW2, self).setModelParameters() # Currently this is the only one with parameters being fitted
+    super(RWWEI2_EEG_BOLD, self).setModelParameters() # Currently this is the only one with parameters being fitted
     self.eeg.setModelParameters()
     self.bold.setModelParameters()    
 
 def createIC(self, ver):
-    #super(mmRWW2, self).createIC()
+    #super(RWWEI2_EEG_BOLD, self).createIC()
     #self.eeg.createIC()
     #self.bold.createIC()
 
@@ -55,7 +55,7 @@ def createIC(self, ver):
 
 def forward(self, external, hx, hE):
 
-    NMM_vals, hE = super(mmRWW2_np, self).forward(external, self.next_start_state[:, 0:2], hE) #TODO: Fix the hx in the future
+    NMM_vals, hE = super(RWWEI2_EEG_BOLD_np, self).forward(external, self.next_start_state[:, 0:2], hE) #TODO: Fix the hx in the future
     EEG_vals, hE = self.eeg.forward(self.step_size, self.sim_len, NMM_vals["E"])
     BOLD_vals, hE = self.bold.forward(self.next_start_state[:, 2:6], self.step_size, self.sim_len, NMM_vals["E"])
 

whobpyt/models/RWW2/ParamsRWW2.py → whobpyt/models/RWWEI2/ParamsRWWEI2.py

@@ -1,7 +1,7 @@
 import torch
 from whobpyt.datatypes import AbstractParams, par
 
-class ParamsRWW2(AbstractParams):
+class ParamsRWWEI2(AbstractParams):
     ## EQUATIONS & BIOLOGICAL VARIABLES FROM:
     #
     # Deco G, Ponce-Alvarez A, Hagmann P, Romani GL, Mantini D, Corbetta M. How local excitation–inhibition ratio impacts the whole brain dynamics. Journal of Neuroscience. 2014 Jun 4;34(23):7886-98.

whobpyt/models/RWW2/README.md → whobpyt/models/RWWEI2/README.md

@@ -1,4 +1,4 @@
-# Reduced Wong Wang Neural Mass Model (Implementation 2)
+# Reduced Wong Wang Excitatory Inhibitory Neural Mass Model (Implementation 2)
 
 ## Description:
 

whobpyt/models/RWW2/RWW2.py → whobpyt/models/RWWEI2/RWWEI2.py

@@ -2,7 +2,7 @@
 from whobpyt.datatypes import AbstractNMM, AbstractParams, par
 from math import sqrt
 
-class RWW2(AbstractNMM):
+class RWWEI2(AbstractNMM):
     '''
     Reduced Wong Wang Excitatory Inhibatory (RWWEXcInh) Model - Version 2
 
@@ -16,7 +16,7 @@ class RWW2(AbstractNMM):
 
     Attributes
     -------------
-    params : ParamsRWW2
+    params : ParamsRWWEI2
         An AbstractParams object which contains the model's parameters
     step_size : Float
         The step size of numerical integration (in msec)
@@ -61,7 +61,7 @@ class RWW2(AbstractNMM):
     def __init__(self, num_regions, params, Con_Mtx, Dist_Mtx, step_size = 0.1, sim_len = 1000, useBC = False, device = torch.device('cpu')):  
         '''
         '''
-        super(RWW2, self).__init__() # To inherit parameters attribute
+        super(RWWEI2, self).__init__() # To inherit parameters attribute
 
         # Initialize the RWW Model 
         #

whobpyt/models/RWW2/RWW2_validate.py → whobpyt/models/RWWEI2/RWWEI2_validate.py

@@ -15,7 +15,7 @@
 import numpy
 from math import sqrt
 
-class RWW2_np():
+class RWWEI2_np():
     def __init__(self, num_regions, params, Con_Mtx, Dist_Mtx, step_size = 0.1):        
 
         # Initialize the RWW Model 

whobpyt/models/RWWEI2/__init__.py

@@ -0,0 +1,7 @@
+from .RWWEI2 import RWWEI2
+from .RWWEI2_validate import RWWEI2_np
+
+from .ParamsRWWEI2 import ParamsRWWEI2
+
+from .Multimodal_RWWEI2 import RWWEI2_EEG_BOLD
+from .Multimodal_RWWEI2_validate import RWWEI2_EEG_BOLD_np

whobpyt/models/__init__.py

@@ -3,4 +3,4 @@
 from . import JansenRit
 from . import Linear
 from . import RWW
-from . import RWW2
+from . import RWWEI2