nntrf 1.0.0__py3-none-any.whl

nntrf/loss.py

@@ -0,0 +1,31 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Thu Dec 10 14:56:39 2020
+
+@author: Jin Dou
+"""
+import torch
+def pearsonCorrLoss(x,y):
+
+    mx = torch.mean(x,1)
+    my = torch.mean(y,1)
+    x = x[:,:,0]
+    y = y[:,:,0]
+    xm, ym = x-mx, y-my
+    r_num = torch.sum(xm * ym,1)
+    r_den = torch.sqrt(torch.sum(xm*xm,1) * torch.sum(ym*ym,1))
+#    print(torch.sum(r_num_sub)/r_den)
+    if(torch.mean(r_num) == 0 and torch.mean(r_den) ==0):
+        raise ValueError("gradient gone\n")
+        r = None
+    else:
+        r = r_num / r_den
+#    return 1 - r**2
+        
+#    print(r)
+    r = torch.mean(r)
+#    loss = torch.exp(r)
+    loss = 1 - r
+
+    return loss
+

nntrf/metrics.py

@@ -0,0 +1,35 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Thu Dec 10 14:55:06 2020
+
+@author: Jin Dou
+"""
+import numpy as np
+# from scipy import stats as spStats
+
+def Pearsonr(x,y):
+    nObs = len(x)
+    sumX = np.sum(x,0)
+    sumY = np.sum(y,0)
+    sdXY = np.sqrt((np.sum(x**2,0) - (sumX**2/nObs)) * (np.sum(y ** 2, 0) - (sumY ** 2)/nObs))
+    
+    r = (np.sum(x*y,0) - (sumX * sumY)/nObs) / sdXY
+    return r
+
+#def Pearsonr(x,y):
+##    x = np.squeeze(x)
+##    y = np.squeeze(y)
+##    print(x.shape)
+##    print(y.shape)
+#    
+#    out = spStats.pearsonr(x,y)
+##    print('correlation',out)
+#    return out
+
+def BatchPearsonr(pred,y):
+    result = list()
+    for i in range(len(pred)):
+        out1 = Pearsonr(pred[i],y[i])
+#        print(out1)
+        result.append(out1)
+    return np.mean(result,0)

nntrf/models/__init__.py

@@ -0,0 +1,3 @@
+from .linear import *
+from .nonlinear import *
+from .composite import *

nntrf/models/composite.py

@@ -0,0 +1,63 @@
+import torch
+
+class TwoMixedTRF(torch.nn.Module):
+        
+    def __init__(
+        self,
+        device,
+        trfs, #list of trf models
+        feats_keys, #list of feat key for each trf in the trfs
+    ):
+        super().__init__()
+        self.trfs:torch.nn.Module = torch.nn.ModuleList([trf for trf in trfs])
+        self.feats_keys = feats_keys
+        self.device = device
+
+    def forward(self,feat_dict:dict, y):
+
+        pred_list = []
+        for trf_index, iTRF in enumerate(self.trfs):
+            feats_key = self.feats_keys[trf_index]
+            feats = []
+            n_dict_feat = 0
+            for feat_key in feats_key:
+                # print(feat_dict.keys())
+                feat = feat_dict[feat_key]
+                if isinstance(feat, dict):
+                    feats.append(feat)
+                    n_dict_feat += 1
+                else:
+                    assert isinstance(feat, torch.Tensor)
+                    feats.append(feat)
+            if n_dict_feat > 0:
+                assert len(feats) == n_dict_feat
+                # concatente
+                if len(feats) == 1:
+                    feats = feats[0]
+                else:
+                    # raise NotImplementedError
+                    timeinfo_0 = feats[0]['timeinfo']
+                    xs = []
+                    for feat in feats:
+                        xs.append(feat['x'])
+                        torch.equal(timeinfo_0, feat['timeinfo'])
+                    xs = torch.cat(xs, dim = -2)
+                    feats = {
+                        'x':xs,
+                        'timeinfo':timeinfo_0
+                    }
+                pred_list.append(iTRF(**feats))
+            else:
+                minLen = min([f.shape[-1] for f in feats])
+                feats = torch.cat([f[...,:minLen] for f in feats], axis = -2)
+                pred_list.append(iTRF(feats))
+            
+
+
+        minLen = min([p.shape[-1] for p in pred_list] + [y.shape[-1]])
+            
+        pred_list = [p[...,:minLen] for p in pred_list]
+        cropedY = y[:,:,:minLen]
+        pred = sum(pred_list)
+        # stop
+        return pred,cropedY

nntrf/models/linear.py

@@ -0,0 +1,269 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Wed Dec  9 10:11:27 2020
+
+@author: Jin Dou
+"""
+
+
+import torch
+import numpy as np
+from ..metrics import Pearsonr, BatchPearsonr
+from ..utils import TensorsToNumpy
+try:
+    from matplotlib import pyplot as plt
+except:
+    plt = None
+
+def msec2Idxs(msecRange,fs):
+    '''
+    convert a millisecond range to a list of sample indexes
+    
+    the left and right ranges will both be included
+    '''
+    assert len(msecRange) == 2
+    
+    tmin = msecRange[0]/1e3
+    tmax = msecRange[1]/1e3
+    return list(range(int(np.floor(tmin*fs)),int(np.ceil(tmax*fs)) + 1))
+
+def Idxs2msec(lags,fs):
+    '''
+    convert a list of sample indexes to a millisecond range
+    
+    the left and right ranges will both be included
+    '''
+    temp = np.array(lags)
+    return list(temp/fs * 1e3)
+
+class LRTRF(torch.nn.Module):
+    '''
+    the TRF implemented with a linear layer and time lag of input 
+    '''
+    # the shape of the input for the forward should be the (nBatch,nTimeSteps,nChannels) 
+    
+    def __init__(self,inDim,outDim,tmin_ms,tmax_ms,fs,bias = True):
+        super().__init__()
+        self.tmin_ms = tmin_ms
+        self.tmax_ms = tmax_ms
+        self.fs = fs
+        self.lagIdxs = msec2Idxs([tmin_ms,tmax_ms],fs)
+        self.lagTimes = Idxs2msec(self.lagIdxs,fs)
+        self.realInDim = len(self.lagIdxs) * inDim
+        self.oDense = torch.nn.Linear(self.realInDim,outDim,bias = bias)
+        self.inDim = inDim
+        self.outDim = outDim
+        
+    def timeLagging(self,tensor):
+        x = tensor
+        nBatch = x.shape[0]
+        batchList = []
+        for batchId in range(nBatch):
+            batch = x[batchId:batchId+1]
+            lagDataList = []
+            for idx,lag in enumerate(self.lagIdxs):
+                # we assume the last second dimension indicates time steps
+                if lag < 0:
+                    temp = torch.nn.functional.pad(batch,((0,0,0,-lag)))
+#                    lagDataList.append(temp[:,-lag:,:])
+                    lagDataList.append((temp.T)[:,-lag:].T)
+                elif lag > 0:
+                    temp = torch.nn.functional.pad(batch,((0,0,lag,0)))
+#                    lagDataList.append(temp[:,0:-lag,:])
+                    lagDataList.append((temp.T)[:,0:-lag].T)
+                else:
+                    lagDataList.append(batch)
+            batchList.append(torch.cat(lagDataList,-1))
+        x3 = torch.cat(batchList,0)
+        return x3
+    
+    def forward(self,x):
+        x = self.timeLagging(x)
+        return self.oDense(x)
+    
+    @property
+    def weights(self):
+        return self.state_dict()['oDense.weight'].cpu().detach().numpy()
+    
+    @property
+    def w(self):
+        '''
+        funtion reproduce the definition of w in mTRF-toolbox
+
+        Returns
+        -------
+        None.
+
+        '''
+        w = self.oDense.weight.T.cpu().detach()
+        w = w.view(len(self.lagIdxs),self.inDim,self.outDim)
+        w = w.permute(1,0,2)
+        w = w.numpy()
+        return w
+    
+    def loadFromMTRFpy(self,w,b,device):
+        #w: (nInChan, nLag, nOutChan)
+        # print(w.shape)
+        w = w * 1/ self.fs
+        b = b * 1/self.fs
+        b = b[0]
+        w = torch.FloatTensor(w).to(device)
+        w = w.permute(1,0,2)
+        w = w.reshape(-1,w.shape[-1]).T
+        b = torch.FloatTensor(b).to(device)
+        with torch.no_grad():
+            self.oDense.weight = torch.nn.Parameter(w)
+            self.oDense.bias = torch.nn.Parameter(b)
+        return self
+
+class CPadOrCrop1D(torch.nn.Module):
+    def __init__(self,tmin_idx,tmax_idx):
+        super().__init__()
+        self.tmin_idx = tmin_idx
+        self.tmax_idx = tmax_idx
+    
+    def forward(self,x):
+        # padding buttom
+        if (self.tmin_idx <= 0):
+            x = torch.nn.functional.pad(x,((0,-self.tmin_idx)))
+        else:
+            x = x[:,:,:-self.tmin_idx]
+        if (self.tmax_idx < 0):
+            x = x[:,:,-self.tmax_idx:]
+        else:
+            x = torch.nn.functional.pad(x,((self.tmax_idx,0)))
+        return x
+            
+
+class CNNTRF(torch.nn.Module):
+    '''
+    the TRF implemented with a convolutional layer,
+        and zero padding of input
+    '''
+      
+    # the shape of the input for the forward should be the (nBatch,nChannels,nTimeSteps,) 
+    # Be care of the calculation of correlation, when using this model,
+    # because the nnTRF.Metrics.Pearsonr treat the input data as the shape of 
+    # (nTimeSteps, nChannels)
+    def __init__(self,inDim,outDim,tmin_ms,tmax_ms,fs,groups = 1,enableBN = False, dilation = 1):
+        super().__init__()
+        self.tmin_ms = tmin_ms
+        self.tmax_ms = tmax_ms
+        self.fs = fs
+        self.lagIdxs = msec2Idxs([tmin_ms,tmax_ms],fs)
+        self.lagTimes = Idxs2msec(self.lagIdxs,fs)
+        self.tmin_idx = self.lagIdxs[0]
+        self.tmax_idx = self.lagIdxs[-1]
+        nLags = len(self.lagTimes)
+        nKernels = (nLags - 1) / dilation + 1
+        assert np.ceil(nKernels) == np.floor(nKernels)
+        nKernels = int(nKernels)
+        self.oCNN = torch.nn.Conv1d(
+            inDim, 
+            outDim, 
+            nKernels,
+            groups = groups,
+            dilation = dilation
+        )
+        self.oPadOrCrop = CPadOrCrop1D(self.tmin_idx,self.tmax_idx)
+        self.groups = groups
+        self.enableBN = enableBN
+        self.oBN = torch.nn.BatchNorm1d(inDim,affine=False,track_running_stats=False)
+        self.dilation = dilation
+        #if both lagMin and lagMax > 0, more complex operation
+        
+    def forward(self,x):
+        if self.enableBN:
+            x = self.oBN(x)
+        x = self.oPadOrCrop(x)
+        x = self.oCNN(x)
+        return x
+    
+    @property
+    def weights(self):
+        '''
+        Returns
+        -------
+        Formatted weights, with timeLag dimension flipped to conform to mTRF
+        [outChannels, inChannels, timeLags]
+        '''
+        return np.flip(self.state_dict()['oCNN.weight'].cpu().detach().numpy(),axis = -1)
+    
+    @property
+    def w(self):
+        '''
+        funtion reproduce the definition of w in mTRF-toolbox
+
+        Returns
+        -------
+        None.
+
+        '''
+        tensor = self.state_dict()['oCNN.weight']
+        tensor = tensor.permute(1,2,0)
+        return np.flip(tensor.cpu().detach().numpy(),axis = 1)
+    
+    
+    @property
+    def b(self):
+        return self.oCNN.bias.squeeze().detach().cpu().numpy()
+    
+    def loadFromMTRFpy(self,w,b,device):
+        #w: (nInChan, nLag, nOutChan)
+        w = w * 1/self.fs
+        b = b * 1/self.fs
+        b = b[0]
+        w = np.flip(w,axis = 1).copy()
+        w = torch.FloatTensor(w).to(device)
+        w = w.permute(2,0,1)
+        b = torch.FloatTensor(b).to(device)
+        with torch.no_grad():
+            self.oCNN.weight = torch.nn.Parameter(w)
+            self.oCNN.bias = torch.nn.Parameter(b)
+        return self
+    
+    @property
+    def t(self):
+        return self.lagTimes
+    
+    def BatchPearsonr(self,pred,y):
+        tensors = TensorsToNumpy(pred.transpose(-1,-2),y.transpose(-1,-2))
+        return BatchPearsonr(*tensors)
+    
+    @property
+    def readableWeights(self):
+        '''
+        Returns
+        -------
+        Readable formatted weights
+        [timeLags, inChannels, outChannels]
+        '''
+        return self.weights.T
+    
+    def W(self,inIdx=None,outIdx=None,tIdx=None):
+        '''
+        Returns
+        -------
+        readable formatted weights of selected inChannel/outChannel/lagTime
+        '''
+        
+        inIdx = slice(inIdx) if inIdx is None else inIdx
+        outIdx = slice(outIdx) if outIdx is None else outIdx
+        tIdx = slice(tIdx) if tIdx is None else tIdx
+        return self.readableWeights[tIdx,inIdx,outIdx]
+    
+    def load(self,path):
+        self.load_state_dict(torch.load(path,map_location='cpu')['state_dict'])
+        self.eval()
+        
+        
+    def plotWeights(self,outChan = None, inChan = None):
+        fig,ax = plt.subplots()
+        if outChan is None:
+            outChan = slice(outChan)
+        if inChan is None:
+            inChan = slice(inChan)
+        ax.plot(self.t[::self.dilation], self.weights[outChan, inChan].T)
+        return fig, ax
+    
+