nntrf 1.0.0__py3-none-any.whl

nntrf/models/nonlinear.py

@@ -0,0 +1,1358 @@
+import math
+import numpy as np
+import torch
+from scipy.stats import pearsonr
+from torch.nn.functional import pad, fold
+from .linear import msec2Idxs, Idxs2msec, CPadOrCrop1D
+try:
+    import skfda
+except:
+    skfda = None
+
+try:
+    from matplotlib import pyplot as plt
+except:
+    plt = None
+
+try:
+    from mtrf.model import TRF
+except:
+    TRF = None
+
+
+def fit_forward_mtrf(stim, resp, fs, tmin_ms, tmax_ms, regularization, k):
+    trf = TRF(direction=1)
+    trf.train(stim, resp, fs, tmin_ms / 1e3, tmax_ms / 1e3, regularization, k = k)
+    return trf.weights, trf.bias
+
+def seqLast_pad_zero(seq, value = 0):
+    maxLen = max([i.shape[-1] for i in seq])
+    output = []
+    for i in seq:
+        output.append(pad(i,(0,maxLen - i.shape[-1]), value = value))
+    return torch.stack(output,0)
+
+class CausalConv(torch.nn.Module):
+
+    def __init__(self,inDim,outDim,nKernel,dilation = 1):
+        super().__init__()
+        self.nKernel = nKernel
+        self.dilation = dilation
+        self.conv = torch.nn.Conv1d(
+            inDim, 
+            outDim, 
+            nKernel, 
+            dilation = dilation
+        )
+    
+    def forward(self,x):
+        '''
+        x: (nBatch, nChan, nSeq)
+        '''
+        # padding left
+        x = torch.nn.functional.pad(x,(( self.dilation * (self.nKernel-1) ,0)))
+
+        #(nBatch, nOutChan, nSeq)
+        x = self.conv(x)
+        return x
+## experiment module end
+    
+class TRFAligner(torch.nn.Module):
+    
+    def __init__(self,device):
+        super().__init__()
+        self.device = device
+
+    def forward(self,TRFs,sourceIdx,nRealLen):#,targetTensor):
+        '''
+        in-place operation
+        Parameters
+        ----------
+        TRFs : TYPE, (nBatch, outDim, nWin, nSeq)
+            tensors output by DyTimeEncoder.
+        sourceIdx : TYPE, (nBatch, nSeq)
+            index of dyImpulse tensor to be assigned to target tensor
+        nRealLen: 
+            the length of the target
+        Returns
+        -------
+        None.
+
+        '''
+        nBatch, outDim, nWin, nSeq = TRFs.shape
+        # (nBatch, outDim, nWin, nSeq)
+        respUnfold = TRFs 
+        maxSrcIdx = torch.max(sourceIdx[:, -1])
+        if maxSrcIdx >= nRealLen:
+            nRealLen = maxSrcIdx + 1
+        # print(outDim,nWin,nRealLen,respUnfold.shape,sourceIdx)
+        self.cache = torch.zeros((nBatch, outDim,nWin,nRealLen),device = self.device)
+
+        idxWin = torch.arange(nWin)
+        idxChan = torch.arange(outDim)
+        idxBatch = torch.arange(nBatch)
+        idxWin = idxWin[:, None]
+        idxChan = idxChan[:,None, None]
+        idxBatch = idxBatch[:, None, None, None]
+        sourceIdx = sourceIdx[:,None, None,:]
+
+        self.cache[idxBatch, idxChan, idxWin, sourceIdx] = respUnfold #(nBatch, outDim,nWin,nRealLen)
+        self.cache = self.cache.view(nBatch,-1,nRealLen) # (nBatch, outDim*nWin, nRealLen)
+        foldOutputSize = (nRealLen + nWin - 1, 1)
+        foldKernelSize = (nWin, 1)
+        #(nBatch,outDim,foldOutputSize,1)
+        output = fold(self.cache,foldOutputSize,foldKernelSize)
+        #(nBatch,outDim,nRealLen)
+        targetTensor = output[:,:,:nRealLen,0]
+        return targetTensor
+
+class LTITRFGen(torch.nn.Module):
+    def __init__(self,inDim,nWin,outDim,ifAddBiasInForward = True):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(outDim,inDim,nWin))
+        self.bias = torch.nn.Parameter(torch.ones(outDim))
+        k = 1 / (inDim * nWin)
+        lower = - np.sqrt(k)
+        upper = np.sqrt(k)
+        torch.nn.init.uniform_(self.weight, a = lower, b = upper)
+        torch.nn.init.uniform_(self.bias, a = lower, b = upper)
+        self.ifAddBiasInForward = ifAddBiasInForward
+
+    @property
+    def outDim(self):
+        return self.weight.shape[0]
+
+    @property
+    def inDim(self):
+        return self.weight.shape[1]
+    
+    @property
+    def nWin(self):
+        return self.weight.shape[2]
+
+    def forward(self,x):
+        # x: (nBatch, inDim, nSeq)
+        assert x.ndim == 3
+        kernelsTemp =  self.weight[None, ..., None] #(1, outDim, inDim, nWin, 1) 
+        xTemp = x[:, None, :, None, :] #(nBatch, 1, inDim, 1, nSeq)
+        TRFs = xTemp * kernelsTemp  #(nBatch, outDim, inDim, nWin, nSeq)
+        if self.ifAddBiasInForward:
+            TRFs = TRFs + self.bias[..., None, None, None] #(nBatch, outDim, inDim, nWin, nSeq)
+        TRFs = TRFs.sum(2) #(nBatch, outDim, nWin, nSeq)
+        return TRFs
+
+    def load_mtrf_weights(self, w, b, fs, device):
+        #w: (nInChan, nLag, nOutChan)
+        b = b[0]
+        w = w * 1 / fs
+        b = b * 1/ fs
+        w = torch.FloatTensor(w).to(device)
+        b = torch.FloatTensor(b).to(device)
+        w = w.permute(2, 0, 1) #(nOutChan, nInChan, nLag)
+        with torch.no_grad():
+            self.weight = torch.nn.Parameter(w)
+            self.bias = torch.nn.Parameter(b)
+        return self
+
+    def export_mtrf_weights(self, fs):
+        with torch.no_grad():
+            # (nInChan, nLag, nOutChan)
+            w = self.weight.cpu().detach().permute(1, 2, 0).numpy()
+            b = self.bias.cpu().detach().numpy().reshape(1,-1)
+        w = w * fs 
+        b = b * fs
+        return w, b
+    
+    def stop_update_weights(self):
+        self.weight.requires_grad_(False)
+        self.weight.grad = None
+        self.bias.requires_grad_(False)
+        self.bias.grad = None
+        
+    def enable_update_weights(self):
+        self.requires_grad_(True)
+        self.weight.grad = torch.zeros_like(
+            self.weight
+        )
+        self.bias.grad = torch.zeros_like(
+            self.bias
+        )
+
+class WordTRFEmbedGenTokenizer():
+    def __init__(self, wordsDict, device):
+        self.wordsDict = wordsDict
+        self.device = device
+
+    def __call__(self, words):
+        batchTokens = []
+        for ws in words:
+            tokens = []
+            for w in ws:
+                tokens.append(self.wordsDict[w])
+            batchTokens.append(
+                torch.tensor(tokens, dtype = torch.long,device = self.device)
+            )
+        return batchTokens
+
+class WordTRFEmbedGen(torch.nn.Module):
+
+    def __init__(
+        self, 
+        outDim,
+        hiddenDim,
+        tmin_ms,
+        tmax_ms,
+        fs,
+        wordsDict,
+        device
+    ):
+        super().__init__()
+        self.outDim = outDim
+        self.hiddenDim = hiddenDim
+        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)
+        nWin = len(self.lagTimes)
+        self.nWin = nWin
+        self.embedding_dim = nWin * hiddenDim
+
+        self.device = device
+        self.wordsDict = wordsDict
+        self.embedding = torch.nn.Embedding(
+            len(wordsDict)+1,
+            self.embedding_dim,
+            padding_idx = 0
+        ).to(device)
+        self.proj = torch.nn.Linear(self.hiddenDim, self.outDim, device = device)
+
+    def forward(self, batchTokens):
+        # (nBatch, outDim, nWin, nSeq)
+        batchTokens = seqLast_pad_zero(batchTokens)
+        # (nBatch, nWin * hiddenDim)
+        trfs = self.embedding(batchTokens)
+        # print(trfs.shape)
+        trfs = trfs.reshape(*trfs.shape[:2], self.hiddenDim, self.nWin)
+        # (nBatch, nSeq, nWin, hiddenDim)
+        # print(trfs.shape)
+        trfs = trfs.permute(0, 1, 3, 2)
+        # (nBatch, nSeq, nWin, outDim)
+        # print(torch.cuda.memory_allocated()/1024/1024)
+        trfs = self.proj(trfs)
+        # (nBatch, outDim, nWin, nSeq)
+        trfs = trfs.permute(0, 3, 2, 1)
+        # print(trfs.shape)
+        return trfs 
+
+
+class CustomKernelCNNTRF(torch.nn.Module):
+
+    def __init__(self,inDim,outDim,tmin_ms,tmax_ms,fs,groups = 1,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)
+        self.nWin = int(nKernels)
+        self.oPadOrCrop = CPadOrCrop1D(self.tmin_idx,self.tmax_idx)
+        self.groups = groups
+        self.dilation = dilation
+        self.inDim = inDim
+        self.outDim = outDim
+
+        k = 1 / (inDim * self.nWin)
+        lower = - np.sqrt(k)
+        upper = np.sqrt(k)
+        self.bias = torch.nn.Parameter(torch.ones(outDim))
+        torch.nn.init.uniform_(self.bias, a = lower, b = upper)
+
+    def setTRFGen(self, trfGen):
+        self.trfGen = trfGen
+
+    def forward(self, x, weight = None):
+        if weight is None:
+            weight = self.trfGen.TRF()
+        assert self.nWin == weight.shape[2]
+        assert self.outDim == weight.shape[0]
+        assert self.inDim == weight.shape[1]
+        x = self.oPadOrCrop(x)
+        y = torch.nn.functional.conv1d(
+            x, 
+            weight, 
+            bias=self.bias, 
+            dilation=self.dilation, 
+            groups=self.groups
+        )
+        return y
+
+class FuncBasisTRF(torch.nn.Module):
+
+    def __init__(self, inDim, outDim, tmin_idx, tmax_idx, timeshiftLimit_idx ,device) -> None:
+        super().__init__()
+        self.timeshiftLimit_idx = torch.tensor(timeshiftLimit_idx)
+        self.time_embedding = self.get_time_embedding(
+            tmin_idx, tmax_idx, device) 
+        self.time_embedding_ext = self.get_time_embedding(
+            tmin_idx-timeshiftLimit_idx, tmax_idx+timeshiftLimit_idx, device) 
+        nWin = self.time_embedding_ext.shape[-2]
+        TRFs = torch.zeros((outDim, inDim, nWin),device=device)
+        self.register_buffer('TRFs',TRFs)
+
+    # def corrected_time_embedding(self, t):
+        # return t + + self.timeshiftLimit_idx
+
+    def TRF(self):
+        # (outDim, inDim, nWin)
+        #(nBatch, 1, 1, nWin, nSeq)
+        return self.forward(1, 0, 1)[0,...,0]#.detach().cpu().numpy()      
+
+    @property
+    def inDim(self):
+        return self.TRFs.shape[1]
+    
+    @property
+    def outDim(self):
+        return self.TRFs.shape[0]
+    
+    @property
+    def nWin(self):
+        return self.TRFs.shape[2]
+    
+    @property
+    def nBasis(self):
+        raise NotImplementedError
+    
+    def vis(self):
+        raise NotImplementedError
+    
+    def forward(self, a,  b, c):
+        raise NotImplementedError
+    
+    def fitTRFs(self, TRFs):
+        raise NotImplementedError
+    
+    def get_time_embedding(self, tmin_idx, tmax_idx, device = 'cpu'):
+        #(1, 1, 1, nWin, 1) 
+        return torch.arange(tmin_idx,tmax_idx+1, device=device)\
+            .view(1,1,1,-1,1)
+    
+    @property
+    def timelag_idx(self,):
+        return self.time_embedding.detach().cpu().squeeze()
+
+    @property
+    def timelag_idx_ext(self,):
+        return self.time_embedding_ext.detach().cpu().squeeze()
+
+
+def build_gaussian_response(x, mu, sigma):
+    # x: (nBatch, 1, 1, nWin, nSeq)
+    # mu: (nBasis)
+    # sigma: (nBasis, outDim, inDim)
+    # output: (nBatch, nBasis, outDim, inDim, nWin, nSeq)
+
+    # x: (nBatch, 1, 1, 1, nWin, nSeq)
+    x = x[:, None, ...]
+    # mu: (nBasis, 1, 1, 1, 1)
+    mu = mu[..., None, None, None, None]
+    # sigma: (nBasis, outDim, inDim,  1, 1)
+    sigma = sigma[..., None, None]
+    # output: (nBatch, nBasis, outDim, inDim, nWin, nSeq)
+    return torch.exp(-(x-mu)**2 / (2*(sigma)**2))
+
+def solve_coef(gaussresps, trf):
+    nWin = trf.shape[0]
+    A = np.concatenate([gaussresps, np.ones((1,nWin))], axis = 0).T
+    coefs = np.linalg.lstsq(A, trf, rcond=None)[0]
+    return coefs
+
+class GaussianBasisTRF(FuncBasisTRF):
+    
+    def __init__(
+        self,
+        inDim,
+        outDim,
+        tmin_idx, 
+        tmax_idx,
+        nBasis,
+        timeshiftLimit_idx = 0,
+        sigmaMin = 6.4,
+        sigmaMax = 6.4,
+        ifSumInDim = False,
+        device = 'cpu',
+        mu = None,
+        sigma = None,
+        include_constant_term = True
+    ):
+        super().__init__(inDim, outDim, tmin_idx, tmax_idx, timeshiftLimit_idx, device)
+        nWin = self.nWin
+        ### Fittable Parameters
+        ## out projection init
+        coefs = torch.ones((nBasis + 1, outDim, inDim), device = device, dtype = torch.float32)
+        torch.nn.init.kaiming_uniform_(coefs, a=math.sqrt(5))
+        self.coefs = torch.nn.Parameter(coefs)
+        ## bias init
+        # k = 1 / (inDim * nWin)
+        # lower = - np.sqrt(k)
+        # upper = np.sqrt(k)
+        # self.bias = torch.nn.Parameter(torch.ones(outDim))
+        # torch.nn.init.uniform_(self.bias, a = lower, b = upper)
+        ## sigma init
+        if sigma is not None:
+            assert len(sigma) == nBasis
+            sigma = torch.tensor(sigma)
+        else:
+            sigma = torch.ones(nBasis, outDim, inDim, device = device, dtype = torch.float32) * (sigmaMin + sigmaMax) / 2
+        self.sigma = torch.nn.Parameter(sigma)
+        # torch.nn.init.uniform_(self.sigma, a = lower, b = upper)
+
+        ### Fixed Values
+        # timeEmbed = torch.arange(nWin) 
+        # self.register_buffer('timeEmbed', timeEmbed)
+        time_embedding_ext = self.time_embedding_ext.squeeze()
+        tmin_idx_ext, tmax_idx_ext = time_embedding_ext[0], time_embedding_ext[-1]
+        if mu is not None:
+            assert len(mu) == nBasis
+            mu = torch.tensor(mu, device = device, dtype = torch.float32)
+        else:
+            mu = torch.linspace(tmin_idx_ext.item(), tmax_idx_ext.item(), nBasis + 2)[1:-1]
+        self.register_buffer('mu', mu)
+        # self.mu = torch.nn.Parameter(mu)
+
+        sigmaMin = torch.tensor(sigmaMin)
+        self.register_buffer('sigmaMin', sigmaMin)
+        sigmaMax = torch.tensor(sigmaMax)
+        self.register_buffer('sigmaMax', sigmaMax)
+        self.ifSumInDim = ifSumInDim
+        self.include_constant_term = include_constant_term
+        self.device = device
+
+    def vec_sum(self, x):
+        return self.vec_gauss_sum(x)
+
+    def vec_gauss_sum(self, x):
+        sigma = self.sigma
+        sigma = torch.maximum(sigma, self.sigmaMin)
+        sigma = torch.minimum(sigma, self.sigmaMax)
+        # print(sigma)
+        # (nBatch, nBasis, outDim, inDim, nWin, nSeq)
+        # print(self.mu, sigma)
+        gaussResps = build_gaussian_response(
+            x, 
+            self.mu, 
+            sigma
+        )
+        
+        # coefs: (nBasis + 1, outDim, inDim, 1, 1)
+        coefs = self.coefs[..., None, None]
+        # print(gaussResps.shape)
+        # nBatch, _, outDim, inDim, nWin, nSeq = gaussResps.shape
+        # (nBatch, nBasis+1, outDim, inDim, nWin, nSeq)
+        # aug_gaussResps = torch.cat([gaussResps, torch.ones(nBatch, 1, outDim, inDim, nWin, nSeq, device = self.device)], dim = -5)
+        # # wGaussResps = coefs[:-1,...] * gaussResps
+        # wGaussResps = coefs * aug_gaussResps
+        # # (nBatch, outDim, inDim, nWin, nSeq)
+        # wGaussResps = wGaussResps.sum((-5))
+
+        # (nBasis, outDim, inDim, 1, 1)
+        coefs_1 = coefs[:-1]
+        # (outDim, inDim, 1, 1)
+        coefs_2 = coefs[-1]
+        # (nBatch, nBasis+1, outDim, inDim, nWin, nSeq)
+        w_gauss_resps_1 = coefs_1 * gaussResps
+        if self.include_constant_term:
+            # print('include constant term')
+            w_gauss_resps = w_gauss_resps_1.sum(-5) +  coefs_2
+        else:
+            w_gauss_resps = w_gauss_resps_1.sum(-5)
+
+        return w_gauss_resps
+
+    def forward(self, a, b, c):
+        # print(x.shape)
+        # output (nBatch, outDim, (inDim), nWin, nSeq)
+        # x: x: (nBatch, 1, 1, nWin, nSeq)
+        # currently just support the 'component' mode
+        x = c * (self.time_embedding - b)
+        # x = self.corrected_time_embedding(x)
+        wGaussResps = a * self.vec_gauss_sum(x)
+        # print(coefs[:,0,0,0,0])
+        if self.ifSumInDim:
+            # (nBatch, outDim, nWin, nSeq)
+            wGaussResps = wGaussResps.sum((-3))
+            # wGaussResps = wGaussResps + self.bias[:, None, None]
+        # print(wGaussResps)
+        return wGaussResps
+    
+    @property
+    def nBasis(self):
+        return self.sigma.shape[0]
+    
+    def fitTRFs(self, TRFs):
+        '''
+        TRFs is the numpy array of mtrf weights
+        Shape: [nInDim, nLags, nOutput]
+
+        self.coefs: (nBasis+1, outDim, inDim)
+        sigma: (nBasis, outDim, inDim)
+        '''
+        # print(TRFs.shape)
+        TRFs = torch.from_numpy(TRFs)
+        TRFs = TRFs.permute(2, 0, 1)
+        self.TRFs[:,:,:] = TRFs.to(self.device)[:,:,:]
+        x = self.time_embedding_ext
+        sigma = self.sigma
+        # print(sigma)
+        # (nBasis, outDim, inDim, nWin)
+        with torch.no_grad():
+            gaussResps = build_gaussian_response(
+                x, 
+                self.mu, 
+                sigma
+            )[0, ..., 0].cpu().numpy()
+
+        nWin = TRFs.shape[2]
+        assert nWin == self.nWin
+        # (nBasis+1, outDim, inDim)
+        coefs = np.zeros(self.coefs.shape)
+        for i in range(self.outDim):
+            for j in range(self.inDim):
+                t_trf = TRFs[i,j,:]
+                # (nBasis, nWin)
+                t_gauss = gaussResps[:, i, j, :] 
+                t_coef = solve_coef(
+                    t_gauss,
+                    t_trf
+                )
+                # print(coefs[:, i, j].shape, t_coef.shape)
+                coefs[:, i, j] = t_coef
+        
+        with torch.no_grad():
+            self.coefs[:,:,:] = torch.from_numpy(coefs)
+            # (nBatch, outDim, inDim, nWin, nSeq)
+            torchTRFs = self.vec_gauss_sum(
+                self.time_embedding_ext,
+            ).cpu().numpy()[0, ..., 0]
+            for j in range(self.outDim):
+                for i in range(self.inDim):
+                    curFTRF = torchTRFs[j, i,:]
+                    TRF = TRFs[j, i, :]
+                    # print(pearsonr(curFTRF, TRF))
+                    assert np.around(pearsonr(curFTRF, TRF)[0]) >= 0.99
+
+    def vis(self, fs = None):
+        if plt is None:
+            raise ValueError('matplotlib should be installed')
+        with torch.no_grad():
+            FTRFs = self.vec_gauss_sum(
+                self.time_embedding_ext,
+            ).cpu()[0, ..., 0]
+        nInChan = self.inDim
+        nOutChan = self.outDim
+        fig, axs = plt.subplots(2)
+        fig.suptitle('top: original TRF, bottom: reconstructed TRF')
+        if fs is None:
+            timelag = self.timelag_idx_ext
+        else:
+            timelag = self.timelag_idx_ext.numpy() / fs
+        for j in range(nOutChan):
+            for i in range(nInChan):
+                TRF = self.TRFs[j,i,:].cpu()
+                FTRF = FTRFs[j,i,:].cpu()
+                # print(pearsonr(FTRF, TRF)[0])
+                axs[0].plot(timelag, TRF)
+                # if j == 0:
+                axs[1].plot(timelag, FTRF)
+        return fig
+
+
+
+class FourierBasisTRF(FuncBasisTRF):
+    
+    def __init__(
+        self,
+        nInChan,
+        nOutChan,
+        tmin_idx,
+        tmax_idx,
+        nBasis,
+        timeshiftLimit_idx,
+        device = 'cpu',
+        if_fit_coefs = False
+    ):
+        #TRFs the TRF for some channels
+        super().__init__(nInChan, nOutChan, tmin_idx, tmax_idx, timeshiftLimit_idx, device)
+        # self.nBasis = nBasis
+        # self.nInChan = nInChan
+        # self.nOutChan = nOutChan
+        # self.nWin = nWin
+        coefs = torch.empty((nOutChan, nInChan, nBasis),device=device)
+        if if_fit_coefs:
+            torch.nn.init.kaiming_uniform_(coefs, a=math.sqrt(5))
+            self.coefs = torch.nn.Parameter(coefs)
+        else:
+            self.register_buffer('coefs', coefs)
+        self.T = self.nWin - 1
+        self.device = device
+        maxN = nBasis // 2
+        self.seqN = torch.arange(1,maxN+1,device = self.device)
+        # self.saveMem = False #expr for saving memory usage
+
+    @property
+    def nBasis(self):
+        return self.coefs.shape[2]
+
+    def fitTRFs(self,TRFs):
+        '''
+        TRFs is the numpy array of mtrf weights
+        Shape: [nInDim, nLags, nOutput]
+        '''
+
+        TRFs = torch.from_numpy(TRFs)
+        TRFs = TRFs.permute(2, 0, 1)
+        self.TRFs[:,:,:] = TRFs.to(self.device)[:,:,:]
+        fd_basis_s = []
+        # grid_points = list(range(self.nWin))
+        grid_points = self.time_embedding_ext.squeeze().cpu().numpy()
+        for j in range(self.outDim):
+            for i in range(self.inDim):
+                TRF = TRFs[j, i, :]
+                fd = skfda.FDataGrid(
+                    data_matrix=TRF,
+                    grid_points=grid_points,
+                )
+                basis = skfda.representation.basis.Fourier(n_basis = self.nBasis)
+                fd_basis = fd.to_basis(basis)
+                coef = fd_basis.coefficients[0]
+                self.coefs[j, i, :] = torch.from_numpy(coef).to(self.device)
+                
+                T = fd_basis.basis.period
+                assert T == self.T
+                fd_basis_s.append(fd_basis)
+                
+        out = self.vec_fourier_sum(
+            self.nBasis,
+            self.T,
+            self.time_embedding_ext,
+            self.coefs
+        )[0, ..., 0]
+        for j in range(self.outDim):
+            for i in range(self.inDim):
+                fd_basis = fd_basis_s[j*self.inDim + i]
+                temp = fd_basis(grid_points).squeeze() #np.arange(0,self.nWin)
+                curFTRF = out[j, i,:].cpu().numpy()
+                TRF = TRFs[j, i,:]
+                assert np.around(pearsonr(TRF, temp)[0]) >= 0.99
+                try:
+                    assert np.allclose(curFTRF,temp,atol = 1e-6)
+                except:
+                    print(TRF, curFTRF,temp)
+                    raise 
+                # print(i,j,pearsonr(TRF, curFTRF))
+                
+    
+    def phi0(self,T):
+        return 1 / ((2 ** 0.5) * ((T/2) ** 0.5))
+
+    def phi2n_1(self,n,T,t):
+        #n: (maxN)
+        #t: (nBatch, 1, 1, nWin, nSeq, 1)
+
+        #(nBatch, 1, 1, nWin, nSeq, maxN)
+        t_input = 2 * torch.pi * t * n / T
+        #(nBatch, 1, 1, nSeq, maxN, nWin)
+        t_input = t_input.permute(0, 1, 2, 4, 5, 3)
+        signal = torch.sin(t_input) / (T/2)**0.5
+        return signal.permute(0, 1, 2, 5, 3, 4)
+    
+    def phi2n(self,n,T,t):
+        #n: (maxN)
+        #t: (nBatch, 1, 1, nWin, nSeq, 1)
+        #(nBatch, 1, 1, nWin, nSeq, maxN)
+        t_input = 2 * torch.pi * t * n / T
+        #(nBatch, 1, 1, nSeq, maxN, nWin)
+        t_input = t_input.permute(0, 1, 2, 4, 5, 3)
+        signal = torch.cos(t_input) / (T/2)**0.5
+        return signal.permute(0, 1, 2, 5, 3, 4)
+    
+    def vec_sum(self, x):
+        return self.vec_fourier_sum(self.nBasis, self.T, x, self.coefs)
+    
+    def vec_fourier_sum(self,nBasis, T, t,coefs):
+        #coefs: (nOutChan, nInChan, nBasis)
+        #t: (nBatch, 1, 1, nWin, nSeq)
+        #   if tChan of t is just 1, which means we share
+        #   the same time-axis transformation for all channels
+        #return: (nBatch, outDim, inDim, nWin, nSeq)
+
+        #(nBatch, 1, 1, nWin, nSeq, 1)
+        t = t[..., None]
+        # print(t.shape)
+        const0 = self.phi0(T)
+        maxN = nBasis // 2
+        # (maxN)
+        seqN = self.seqN
+        # (nBatch, 1, 1, nWin, nSeq, maxN)
+        constSin = self.phi2n_1(seqN, T, t) 
+        constCos = self.phi2n(seqN, T, t)
+
+        # (nBatch, 1, 1, nWin, nSeq, 2 * maxN)
+        constN = torch.stack(
+            [constSin,constCos],
+            axis = -1
+        ).reshape(*constSin.shape[:5], 2*maxN)
+        # print(const0,[i.shape for i in [constN, coefs]])
+
+        nBatch, _, _, nWin, nSeq, nBasis =  constN.shape
+        nOutChan, nInChan, nBasis = coefs.shape
+
+        #(nOutChan, nInChan, 1, 1, nBasis)
+        coefs = coefs[:, :, None, None, :]
+        nBasis = nBasis
+        # print(constN.shape, coefs.shape)
+        '''
+        #expr for saving memory usage
+        memAvai,_ = torch.cuda.mem_get_info()
+        nMemReq = nBatch * nSeq * nInChan * nOutChan * nBasis * nWin * 4 # 4 indicates 4 bytes
+        # print(torch.cuda.memory_allocated()/1024/1024)
+        if nMemReq > memAvai * 0.9 or self.saveMem:
+            out = const0 * coefs[...,0] #(nOutChan, nInChan, 1, 1)
+            for nB in range(2 * maxN):
+                out = out + constN[...,nB] * coefs[...,1+nB]
+        else:
+            # (nbatch, nOutChan, nInChan, nWin, nSeq, nBasis)
+            out =  const0 * coefs[...,0] + (constN * coefs[...,1:]).sum(-1)
+        '''   
+        out =  const0 * coefs[...,0] + (constN * coefs[...,1:]).sum(-1)
+        # print(torch.cuda.memory_allocated()/1024/1024)
+        # (nBatch, outDim, inDim, nWin, nSeq)
+        return out
+    
+    def forward(self,a, b, c):
+        # a,b,c in most strict case (nBatch, 1, 1, 1, nSeq)
+        # loosly: a,b,c can be (nBatch, nOut, nIn, nWin, nSeq)
+
+        #self.time_embedding #(1, 1, 1, nWin, 1) 
+        
+        # x: (nBatch, 1, 1, nWin, nSeq) 
+        nSeq = self.time_embedding
+        x = c * (nSeq - b)
+        # x = self.corrected_time_embedding(x)
+        #(nBatch, outDim, inDim, nWin, nSeq)
+        # nonLinTRFs = aSeq * self.basisTRF( cSeq * ( nSeq -  bSeq) ) 
+
+        # return: 
+        coefs = self.coefs
+        out = a * self.vec_fourier_sum(self.nBasis,self.T,x,coefs)
+        return out
+        
+    def vis(self ,fs = None):
+        if fs is None:
+            timelag = self.timelag_idx_ext
+        else:
+            timelag = self.timelag_idx_ext.numpy() / fs
+        if plt is None:
+            raise ValueError('matplotlib should be installed')
+        with torch.no_grad():
+            FTRFs = self.vec_fourier_sum(
+                self.nBasis,
+                self.T,
+                self.time_embedding_ext,
+                self.coefs
+            )[0, ..., 0]
+        nInChan = self.inDim
+        nOutChan = self.outDim
+        fig, axs = plt.subplots(2)
+        fig.suptitle('top: original TRF, bottom: reconstructed TRF')
+        for j in range(nOutChan):
+            for i in range(nInChan):
+                TRF = self.TRFs[j,i,:].cpu()
+                FTRF = FTRFs[j,i,:].cpu()
+                axs[0].plot(timelag, TRF)
+                axs[1].plot(timelag, FTRF)
+        return fig
+
+basisTRFNameMap = {
+    'gauss': GaussianBasisTRF,
+    'fourier': FourierBasisTRF
+}
+
+
+class TRFsGen(torch.nn.Module):
+
+    def forward(self, x, featOnsetIdx):
+        '''
+        input:
+             x: input to be used to derive the 
+                transformation parameters TRFs
+             featOnsetIdx: time index of item 
+                in the transformed x to be picked as real transformation parameters
+        '''
+        pass
+
+
+class FuncTRFsGen(torch.nn.Module):
+    '''
+    Implement the functional TRF generator, generate dynamically 
+        warped TRF by transform the functional TRF template
+    '''
+
+    def __init__(
+        self, 
+        inDim, 
+        outDim,
+        tmin_ms,
+        tmax_ms,
+        fs,
+        basisTRFName = 'fourier',
+        limitOfShift_idx = 7,
+        nBasis = 21, 
+        mode = '',
+        transformer = None,
+        device = 'cpu',
+        # if_trans_per_outChan = False
+    ):
+        super().__init__()
+        assert mode.replace('+-','') in ['','a','b','a,b','a,b,c','a,c']
+        self.fs = fs
+        self.lagIdxs = msec2Idxs([tmin_ms,tmax_ms],fs)
+        self.lagIdxs_ts = torch.Tensor(self.lagIdxs).float().to(device)
+        self.lagTimes = Idxs2msec(self.lagIdxs,fs)
+        nWin = len(self.lagTimes)
+        self.mode = mode
+        self.transformer:torch.nn.Module = transformer
+        self.n_transform_params = len(mode.split(','))
+        self.device = device
+        if isinstance(basisTRFName, str):
+            self.basisTRF:FuncBasisTRF = basisTRFNameMap[basisTRFName](
+                inDim, 
+                outDim, 
+                self.lagIdxs[0],
+                self.lagIdxs[-1], 
+                nBasis,
+                timeshiftLimit_idx = limitOfShift_idx, 
+                device=device
+            )
+        elif isinstance(basisTRFName, torch.nn.Module):
+            self.basisTRF:FuncBasisTRF = basisTRFName
+        else:
+            raise ValueError()
+        # self.if_trans_per_outChan = if_trans_per_outChan
+
+        if transformer is None:
+            transInDim, transOutDim, device = self.get_default_transformer_param()
+            self.transformer:torch.nn.Module = CausalConv(transInDim, transOutDim, 2).to(device)
+
+        self.limitOfShift_idx = torch.tensor(limitOfShift_idx)
+
+    @classmethod
+    def parse_trans_params(cls,mode):
+        return mode.split(',')
+
+    @property
+    def tmin_ms(self):
+        return self.lagTimes[0]
+    
+    @property
+    def tmax_ms(self):
+        return self.lagTimes[-1]
+
+    @property
+    def inDim(self):
+        return self.basisTRF.inDim
+    
+    @property
+    def outDim(self):
+        return self.basisTRF.outDim
+    
+    @property
+    def nWin(self):
+        return self.basisTRF.nWin - 2 * self.limitOfShift_idx
+    
+    @property
+    def nBasis(self):
+        return self.basisTRF.nBasis
+
+    @property
+    def extendedTimeLagRange(self):
+        minLagIdx = self.lagIdxs[0]
+        maxLagIdx = self.lagIdxs[-1]
+        left = np.arange(minLagIdx - self.limitOfShift_idx, minLagIdx)
+        right = np.arange(maxLagIdx + 1, maxLagIdx + 1 + self.limitOfShift_idx)
+        extLag_idx = np.concatenate([left, self.lagIdxs, right])
+        try:
+            assert len(extLag_idx) == self.nWin + 2 * self.limitOfShift_idx
+        except:
+            print(len(extLag_idx), self.nWin + 2 * self.limitOfShift_idx)
+        timelags = Idxs2msec(extLag_idx, self.fs)
+        return timelags[0], timelags[-1]
+
+    def get_default_transformer_param(self):
+        inDim = self.inDim
+        device = self.device
+        outDim = self.n_transform_params
+        return inDim, outDim, device
+
+    def fitFuncTRF(self, w):
+        w = w * 1 / self.fs
+        with torch.no_grad():
+            self.basisTRF.fitTRFs(w)
+        return self
+    
+    def pickParam(self,paramSeqs,idx):
+        #paramSeqs: (nBatch, nMiddleParam, nOut/1, 1, 1, nSeq)
+        return paramSeqs[:, idx, ...]
+    
+    
+    def getTransformParams(self, x, startIdx = None):
+        paramSeqs = self.transformer(x) #(nBatch, nMiddleParam, nSeq)
+        nBatch, nMiddleParam, nSeq = paramSeqs.shape
+        if startIdx is not None:
+            idxBatch = torch.arange(nBatch)
+            idxMiddleParam = torch.arange(nMiddleParam)
+            idxMiddleParam = idxMiddleParam[:, None]
+            startIdx = startIdx[:, None, :]
+            idxBatch = idxBatch[:, None, None]
+            paramSeqs = paramSeqs[idxBatch, idxMiddleParam, startIdx]
+        
+        #(nBatch, n_transform_params, 1, 1, nSeq), this is the most strict case
+        #however, it can also be (nBatch, n_transform_params, nOut, 1, nSeq) 
+        #  for different transformation for different channel
+            
+        paramSeqs = paramSeqs.view(nBatch, self.n_transform_params, -1, 1, 1, nSeq) #[:, :, None, None, :]
+        if paramSeqs.shape[2] != 1:
+            assert paramSeqs.shape[2] == self.outDim
+        midParamList = self.mode.split(',')
+        if midParamList == ['']:
+            midParamList = []
+        nParamMiss = 0
+        if 'a' in midParamList:
+            aIdx = midParamList.index('a')
+            #(nBatch, 1, nOut/1, 1, nSeq)
+            aSeq = self.pickParam(paramSeqs, aIdx) 
+            aSeq = torch.abs(aSeq)
+        elif '+-a' in midParamList:
+            aIdx = midParamList.index('+-a')
+            #(nBatch, 1, nOut/1, 1, nSeq)
+            aSeq = self.pickParam(paramSeqs, aIdx)
+        else:
+            nParamMiss += 1
+            #(nBatch, 1, inDim, 1, nSeq)
+            aSeq = x[:, None, :, None, :]
+        if 'b' in midParamList:
+            bIdx = midParamList.index('b')
+            #(nBatch, 1, 1, 1, nSeq)
+            bSeq = self.pickParam(paramSeqs, bIdx) 
+            bSeq = torch.maximum(bSeq, - self.limitOfShift_idx)
+            bSeq = torch.minimum(bSeq,   self.limitOfShift_idx)
+        else:
+            nParamMiss += 1
+            bSeq = 0
+            
+        if 'c' in midParamList:
+            cIdx = midParamList.index('c')
+            #(nBatch, 1, 1, 1, nSeq)
+            cSeq = self.pickParam(paramSeqs, cIdx)
+            #two reasons, cSeq must be larger than 0; 
+            #if 1 is the optimum, abs will have two x for the optimum, 
+            # which is not stable 
+            cSeq =  1 + cSeq
+            cSeq = torch.maximum(cSeq, torch.tensor(0.5))
+            cSeq = torch.minimum(cSeq, torch.tensor(1.28))
+        else:
+            nParamMiss += 1
+            cSeq = 1
+
+        assert (len(midParamList) + nParamMiss) == 3
+        return aSeq, bSeq, cSeq
+
+    def forward(self, x, featOnsetIdx = None):
+        '''
+        x: (nBatch, inDim, nSeq)
+        output: TRFs (nBatch, outDim, nWin, nSeq)
+        '''
+        #(nBatch, nOut/1, 1, 1, nSeq)
+        aSeq, bSeq, cSeq = self.getTransformParams(x, featOnsetIdx)
+        #(1, 1, 1, nWin, 1) 
+        # nSeq = self.lagIdxs_ts[None, None, None, :, None] + self.limitOfShift_idx 
+        # print(aSeq, bSeq, cSeq)x
+        #(nBatch, outDim, inDim, nWin, nSeq)
+        # nonLinTRFs = aSeq * self.basisTRF( cSeq * ( nSeq -  bSeq) ) 
+        # print(aSeq.shape, bSeq.shape)
+        nonLinTRFs = self.basisTRF(aSeq, bSeq, cSeq)
+        # print(torch.cuda.memory_allocated()/1024/1024)
+
+        #(nBatch, outDim, nWin, nSeq)
+        TRFs = nonLinTRFs.sum(2)
+        # print(torch.cuda.memory_allocated()/1024/1024)
+        return TRFs
+
+class ASTRF(torch.nn.Module):
+    '''
+    the TRF implemented the convolution sum of temporal response,
+        (i.e., time-aligning the temporal responses at their 
+        corresponding location, and point-wise sum them).
+        It requres a module to generate temproal responses to each
+        individual stimuli, and also require time information to
+        displace/align the temporal responses at the right
+        indices/location 
+
+    limitation: can't do TRF for zscored input, under this condition
+      location with no stimulus will be non-zero.
+
+
+    the core mechanism of this module following thess steps:
+        1. generate TRFs using the input param 'x',
+        2. determine the onset time of these TRFs within the output time series, using input param 'timeinfo',
+        3. align and sum the generated TRFs at their corresponding time location in the output.
+    
+    Note:
+        1. currently, there are two types of x supported, 
+            type 1: discrete type, means there is a single x for each time point in the 'timestamp'
+            type 2: continuous type, means there is a timeseries of x which has the same length as the output timeseries,
+                and when generating TRFs, only part of the x will actually contribute to this process.
+
+    '''
+
+    def __init__(
+        self,
+        inDim,
+        outDim,
+        tmin_ms,
+        tmax_ms,
+        fs,
+        trfsGen = None,
+        device = 'cpu',
+        x_is_timeseries = False
+    ):
+        '''
+        inDim: int, the number of columns of input 
+        outDim: int, the number of columns of output of ltiTRFGen and trfsGen
+
+        '''
+        super().__init__()
+        assert tmin_ms >= 0
+        self.x_is_timeseries = x_is_timeseries
+        self.lagIdxs = msec2Idxs([tmin_ms,tmax_ms],fs)
+        self.lagTimes = Idxs2msec(self.lagIdxs,fs)
+        nWin = len(self.lagTimes)
+        self.ltiTRFsGen = LTITRFGen(
+            inDim,
+            nWin,
+            outDim,
+            ifAddBiasInForward=False
+        ).to(device)
+        # if callable(trfsGen):
+        #     trfsGen = trfsGen()
+        self.trfsGen:FuncTRFsGen = trfsGen if trfsGen is None else trfsGen.to(device)
+        self.fs = fs
+
+        self.bias = None
+        #also train bias for the trfsGen provided by the user
+        if self.trfsGen is not None:
+            self.init_nonLinTRFs_bias(inDim, nWin, outDim, device)
+        
+        self.trfAligner = TRFAligner(device)
+        self._enableUserTRFGen = False 
+        self.device = device
+
+    @property
+    def inDim(self):
+        return self.ltiTRFsGen.inDim
+    
+    @property
+    def outDim(self):
+        return self.ltiTRFsGen.outDim
+    
+    @property
+    def nWin(self):
+        return self.ltiTRFsGen.nWin
+    
+    @property
+    def tmin_ms(self):
+        return self.lagTimes[0]
+    
+    @property
+    def tmax_ms(self):
+        return self.lagTimes[-1]
+    
+    def init_nonLinTRFs_bias(self, inDim, nWin, outDim, device):
+        self.bias = torch.nn.Parameter(torch.ones(outDim))
+        fan_in = inDim * nWin
+        bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+        torch.nn.init.uniform_(self.bias, -bound, bound)
+
+    def set_trfs_gen(self, trfsGen):
+        self.trfsGen = trfsGen.to(self.device)
+        self.bias = torch.nn.Parameter(
+            torch.ones(self.outDim, device = self.device)
+        )
+        fan_in = self.inDim * self.nWin
+        bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+        torch.nn.init.uniform_(self.bias, -bound, bound)
+
+    def get_params_for_train(self):
+        return [i for i in self.trfsGen.transformer.parameters()] + [self.bias]
+        # raise NotImplementedError()
+        
+
+    def set_linear_weights(self, w, b):
+        #w: (nInChan, nLag, nOutChan)
+        self.ltiTRFsGen.load_mtrf_weights(
+            w,
+            b,
+            self.fs,
+            self.device
+        )
+        return self
+
+    def get_linear_weights(self):
+        w, b = self.ltiTRFsGen.export_mtrf_weights(
+            self.fs
+        )
+        return w, b
+
+    @property
+    def if_enable_trfsGen(self):
+        return self._enableUserTRFGen
+
+    @if_enable_trfsGen.setter
+    def if_enable_trfsGen(self,x):
+        assert isinstance(x, bool)
+        print('set ifEnableNonLin',x)
+        if x == True and self.trfsGen is None:
+            raise ValueError('trfGen is None, cannot be enabled')
+        self._enableUserTRFGen = x
+
+    def stop_update_linear(self):
+        self.ltiTRFsGen.stop_update_weights()
+        
+    def enable_update_linear(self):
+        self.ltiTRFsGen.enable_update_weights()
+
+    def forward(self, x, timeinfo):
+        '''
+        input: 
+            x: nBatch * [nChan, nSeq], 
+            timeinfo: nBatch * [2, nSeq]
+        output: targetTensor
+        '''
+
+        ### record the necessary information of each item in the batch
+        ### for x and targetTensor
+        nSeqs = [] # length of x in the batch
+        nRealLens = [] # length of real output in the batch
+        trfOnsetIdxs = [] # the corresponding index of timepoint in the timeinfo in the batch
+        for ix, xi in enumerate(x):
+            nLenXi = xi.shape[-1]
+            if timeinfo[ix] is not None:
+                # print(timeinfo[ix].shape)
+                if not self.x_is_timeseries:
+                    assert timeinfo[ix].shape[-1] == xi.shape[-1]
+                nLen = torch.ceil(
+                    timeinfo[ix][0][-1] * self.fs
+                ).long() + self.nWin
+                onsetIdx = torch.round(
+                    timeinfo[ix][0,:] * self.fs
+                ).long() + self.lagIdxs[0]
+            else:
+                nLen = nLenXi
+                onsetIdx = torch.tensor(np.arange(nLen)) + self.lagIdxs[0]
+            nSeqs.append(nLenXi)
+            nRealLens.append(nLen)
+            trfOnsetIdxs.append(onsetIdx)
+
+        nGlobLen = max(nRealLens)
+        x = seqLast_pad_zero(x)
+        trfOnsetIdxs = seqLast_pad_zero(trfOnsetIdxs, value = -1)
+        
+        # if x is time series
+        featOnsetIdxs = None
+        if self.x_is_timeseries:
+            # featIdxs is the index where we get the feat for generating TRFs
+            featOnsetIdxs = trfOnsetIdxs.detach().clone()
+            featOnsetIdxs[featOnsetIdxs != -1] =\
+                featOnsetIdxs[featOnsetIdxs != -1] - self.lagIdxs[0]
+        
+        #TRFs shape: (nBatch, outDim, nWin, nSeq)
+        # print(x.shape)
+        TRFs = self.get_trfs(x, featOnsetIdxs)
+
+        #targetTensor shape: (nBatch,outDim,nRealLen)
+        targetTensor = self.trfAligner(TRFs,trfOnsetIdxs,nGlobLen)
+
+        if self.if_enable_trfsGen:
+            targetTensor = targetTensor + self.bias.view(-1,1)
+        else:
+            ltiTRFBias = self.ltiTRFsGen.bias
+            targetTensor = targetTensor + ltiTRFBias.view(-1,1)
+
+        return targetTensor
+    
+    def get_trfs(self, x, featOnsetIdxs = None):
+        if self.if_enable_trfsGen:
+            return self.trfsGen(x, featOnsetIdxs)
+        else:
+            return self.ltiTRFsGen(x)
+        
+
+class ASCNNTRF(ASTRF):
+    #perform CNNTRF within intervals,
+    #and change the weights
+
+    def __init__(
+        self,
+        inDim,
+        outDim,
+        tmin_ms,
+        tmax_ms,
+        fs,
+        trfsGen = None,
+        device = 'cpu'
+    ):
+        torch.nn.Module.__init__(self)
+        # assert tmin_ms >= 0
+        self.inDim = inDim
+        self.outDim = outDim
+        self.tmin_ms = tmin_ms
+        self.tmax_ms = tmax_ms
+        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]
+        nWin = len(self.lagTimes)
+        self._nWin = nWin
+        self.ltiTRFsGen = LTITRFGen(
+            inDim,
+            nWin,
+            outDim,
+            ifAddBiasInForward=False
+        ).to(device)
+        self.trfsGen = trfsGen if trfsGen is None else trfsGen.to(device)
+        self.fs = fs
+
+        self.bias = None
+        #also train bias for the trfsGen provided by the user
+        if self.trfsGen is not None:
+            self.bias = torch.nn.Parameter(torch.ones(outDim, device = device))
+            fan_in = inDim * nWin
+            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+            torch.nn.init.uniform_(self.bias, -bound, bound)
+        
+        self._enableUserTRFGen = False 
+        self.device = device
+
+    @property
+    def nWin(self):
+        return self._nWin
+    
+    def getTRFs(self, ctx):
+        # ctx: (nBatch, inDim,nSeq)
+        # #nTRFs is the number of TRFs needed
+
+        #note the difference between the ASTRF and ASCNNTRF at here
+        #the LTITRF is not multiplied with x
+
+        
+        if self.ifEnableUserTRFGen:
+            #how to decide how much trfs to return????
+            TRFs = self.trfsGen(ctx)
+            #TRFs (nBatch, outDim, nWin, nSeq)
+            TRFs = TRFs[0]  #(outDim, nWin, nSeq)
+            TRFs = TRFs.permute(2, 0, 1)[..., None, :]
+        else:
+            #ctx: (nBatch, inDim,nSeq) (nTRFs, inDim,nSeq)
+            nTRFs = ctx.shape[-1]#len(ctx)
+            # TRFs nTRFs * (nChanOut, nChanIn, nWin)
+            TRFs = [self.ltiTRFsGen.weight] * nTRFs
+        return TRFs
+
+    def getTRFSwitchOnsets(self, x):
+        #X: nBatch * [nChan, nSeq]
+        return [0, 200, 300, 400, 500, 600, 800, 1000]
+
+
+    def defaultCtx(self, switchOnsets, x):
+        switchOnsets2 = switchOnsets + [-1]
+        ctx = []
+        for i in range(len(switchOnsets)):
+            ctx.append(x[:, :, switchOnsets2[i]:switchOnsets2[i+1]])
+        return ctx
+
+    def forward(self, x, timeInfo = None, ctx = None):
+        
+        #currently only support single batch
+
+        #X: nBatch * [nChan, nSeq]
+        #timeinfo: nBatch * [2, nSeq]
+        if timeInfo is None:
+            TRFSwitchOnsets = self.getTRFSwitchOnsets(x)
+        else:
+            TRFSwitchOnsets = timeInfo
+        if ctx is None:
+            ctx = self.defaultCtx(TRFSwitchOnsets, x)
+        nTRFs = len(TRFSwitchOnsets)
+        nBatch, nChan, nSeq = x.shape
+        TRFs = self.getTRFs(ctx)
+        TRFsFlip = [TRF.flip([-1]) for TRF in TRFs]
+        # print([torch.equal(TRFsFlip[0], temp) for temp in TRFsFlip[1:]])
+        # print(TRFsFlip)
+        TRFSwitchOnsets.append(None)
+
+        nPaddedOutput = nSeq \
+            + max(-self.tmin_idx, 0) \
+            + max( self.tmax_idx, 0)
+
+        output = torch.zeros(
+            nBatch, 
+            self.outDim, 
+            nPaddedOutput, 
+            device = self.device
+        )
+
+        #segment startIdx offset
+        startOffset = self.tmin_idx
+        #global startIdx offset
+        offset = min(0, self.tmin_idx)
+
+        realOffset = startOffset - offset
+
+        for idx, TRFFlip in enumerate(TRFsFlip):
+            t_start = TRFSwitchOnsets[idx]
+            t_end = TRFSwitchOnsets[idx+1]
+            t_x = x[..., t_start : t_end]
+            segment = self.trfCNN(t_x,TRFFlip)
+            # print(realOffset)
+            t_startReal = t_start + realOffset
+            t_endReal = t_startReal + segment.shape[-1]
+            # print(segment.shape, t_startReal, t_endReal, t_x.shape, x.shape)
+            output[:,:,t_startReal : t_endReal] += segment 
+
+       #decide how to crop the output based on tmin and tmax
+        startIdx = max(-self.tmin_idx, 0)
+        lenOutput = output.shape[-1]
+        endIdx   = lenOutput - max( self.tmax_idx, 0)
+        
+        return output[..., startIdx: endIdx] + self.ltiTRFsGen.bias.view(-1, 1)
+
+    def trfCNN(self, x, TRFFlip):
+        #X: (nBatch, nChan, nSubSeq)
+        # TRFsFlip is the TRFs with its kernel dimension flipped 
+        # for Conv1D!
+        
+        #need to first padding
+        #timelag doesn't influence how much to pad
+        #  but the offset of the startIdx
+        nWin = TRFFlip.shape[-1]
+        x = torch.nn.functional.pad(x, (nWin-1, nWin-1))
+        #then do the conv
+        output = torch.nn.functional.conv1d(x, TRFFlip)
+        return output
+