text:
# Copyright (c) Hu Zhiming 2020/6/1 jimmyhu@pku.edu.cn All Rights Reserved.

import torch
import torch.nn as nn
#from torchvision import models
import torch.nn.init as init
import torch.nn.functional as F
from math import floor
import numpy as np

# Device configuration
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

class FixationNet(nn.Module):
    def __init__(self, gazeSeqSize, headSeqSize, objectSeqSize, saliencySize, clusterPath):
        super().__init__()
        # the input params
        self.gazeSeqSize = gazeSeqSize
        self.headSeqSize = headSeqSize
        self.objectSeqSize = objectSeqSize
        self.saliencySize = saliencySize

        print('gazeSeqSize: {}'.format(self.gazeSeqSize))
        print('headSeqSize: {}'.format(self.headSeqSize))
        print('taskObjectSeqSize: {}'.format(self.objectSeqSize))
        print('saliencySize: {}'.format(self.saliencySize))

        # preset params
        self.gazeFeatureNum = 2
        self.gazeSeqLength = int(self.gazeSeqSize/self.gazeFeatureNum)
        self.headFeatureNum = 2
        self.headSeqLength = int(self.headSeqSize/self.headFeatureNum)
        # the nearest 3 task-related objects, 3*4
        self.objectFeatureNum = 12
        self.objectSeqLength = int(self.objectSeqSize/self.objectFeatureNum)
        self.saliencyWidth = 24
        self.saliencyNum = int(self.saliencySize/(self.saliencyWidth*self.saliencyWidth))

        # load the cluster centers
        self.cluster = torch.from_numpy(np.load(clusterPath)).float().to(device)
        self.clusterSize = self.cluster.shape[0]
        print('Cluster Size: {}'.format(self.clusterSize))

        # GazeSeqCNN1D Module
        gazeSeqCNN1D_outChannels1 = 32
        gazeSeqCNN1D_poolingRate1 = 2
        gazeSeqCNN1D_kernelSize1 = 1
        gazeSeqCNN1D_featureSize1 = floor((self.gazeSeqLength - gazeSeqCNN1D_kernelSize1 + 1)/gazeSeqCNN1D_poolingRate1)
        self.gazeSeqCNN1D_outputSize = gazeSeqCNN1D_featureSize1 * gazeSeqCNN1D_outChannels1
        self.gazeSeqCNN1D = nn.Sequential(
            nn.Conv1d(in_channels=self.gazeFeatureNum, out_channels=gazeSeqCNN1D_outChannels1,kernel_size=gazeSeqCNN1D_kernelSize1),
            nn.BatchNorm1d(gazeSeqCNN1D_outChannels1),
            nn.ReLU(),
            #nn.LeakyReLU(),
            #nn.Sigmoid(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.MaxPool1d(gazeSeqCNN1D_poolingRate1),
             )

        gazeSeqCNN_dropoutRate = 0.0
        self.gazeSeqCNNDropout = nn.Sequential(
            nn.Dropout(p = gazeSeqCNN_dropoutRate)
             )

        # HeadSeqCNN1D Module
        headSeqCNN1D_outChannels1 = 64
        headSeqCNN1D_poolingRate1 = 2
        headSeqCNN1D_kernelSize1 = 1
        headSeqCNN1D_featureSize1 = floor((self.headSeqLength - headSeqCNN1D_kernelSize1 + 1)/headSeqCNN1D_poolingRate1)
        self.headSeqCNN1D_outputSize = headSeqCNN1D_featureSize1 * headSeqCNN1D_outChannels1
        self.headSeqCNN1D = nn.Sequential(
            nn.Conv1d(in_channels=self.headFeatureNum, out_channels=headSeqCNN1D_outChannels1,kernel_size=headSeqCNN1D_kernelSize1),
            nn.BatchNorm1d(headSeqCNN1D_outChannels1),
            nn.ReLU(),
            #nn.LeakyReLU(),
            #nn.Sigmoid(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.MaxPool1d(headSeqCNN1D_poolingRate1),
            #nn.Conv1d(in_channels=headSeqCNN1D_outChannels1, out_channels=headSeqCNN1D_outChannels2,kernel_size=headSeqCNN1D_kernelSize2),
            #nn.BatchNorm1d(headSeqCNN1D_outChannels2),
            #nn.Sigmoid(),
            #nn.MaxPool1d(headSeqCNN1D_poolingRate2),
             )

        headSeqCNN_dropoutRate = 0.0
        self.headSeqCNNDropout = nn.Sequential(
            nn.Dropout(p = headSeqCNN_dropoutRate)
             )

        # objectSeqCNN1D Module
        objectSeqCNN1D_outChannels1 = 64
        objectSeqCNN1D_poolingRate1 = 2
        objectSeqCNN1D_kernelSize1 = 1
        objectSeqCNN1D_featureSize1 = floor((self.objectSeqLength - objectSeqCNN1D_kernelSize1 + 1)/objectSeqCNN1D_poolingRate1)
        objectSeqCNN1D_outChannels2 = 32
        objectSeqCNN1D_poolingRate2 = 2
        objectSeqCNN1D_kernelSize2 = 1
        objectSeqCNN1D_featureSize2 = floor((objectSeqCNN1D_featureSize1 - objectSeqCNN1D_kernelSize2 + 1)/objectSeqCNN1D_poolingRate2)
        self.objectSeqCNN1D_outputSize = objectSeqCNN1D_featureSize2 * objectSeqCNN1D_outChannels2
        self.objectSeqCNN1D = nn.Sequential(
            nn.Conv1d(in_channels=self.objectFeatureNum, out_channels=objectSeqCNN1D_outChannels1,kernel_size=objectSeqCNN1D_kernelSize1),
            nn.BatchNorm1d(objectSeqCNN1D_outChannels1),
            nn.ReLU(),
            #nn.Sigmoid(),
            #nn.LeakyReLU(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.MaxPool1d(objectSeqCNN1D_poolingRate1),
            nn.Conv1d(in_channels=objectSeqCNN1D_outChannels1, out_channels=objectSeqCNN1D_outChannels2,kernel_size=objectSeqCNN1D_kernelSize2),
            nn.BatchNorm1d(objectSeqCNN1D_outChannels2),
            nn.ReLU(),
            #nn.Sigmoid(),
            nn.MaxPool1d(objectSeqCNN1D_poolingRate2),
             )

        objectSeqCNN_dropoutRate = 0.0
        self.objectSeqCNNDropout = nn.Sequential(
            nn.Dropout(p = objectSeqCNN_dropoutRate)
             )

        # SalCNN Module
        salCNN_outChannels1 = 8
        salCNN_poolingRate1 = 2
        salCNN_kernelSize1 = 1
        salCNN_padding1 = int((salCNN_kernelSize1-1)/2)
        salCNN_imageSize1 = floor((self.saliencyWidth - salCNN_kernelSize1 + 2*salCNN_padding1 + 1)/salCNN_poolingRate1)
        self.salCNN_outputSize = salCNN_imageSize1 * salCNN_imageSize1 * salCNN_outChannels1
        self.SalCNN = nn.Sequential(
            nn.Conv2d(self.saliencyNum, salCNN_outChannels1, kernel_size=salCNN_kernelSize1, padding=salCNN_padding1),
            nn.BatchNorm2d(salCNN_outChannels1),
            nn.ReLU(),
            #nn.Sigmoid(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.MaxPool2d(kernel_size=salCNN_poolingRate1),
             )

        salCNN_dropoutRate = 0.5
        self.SalCNNDropout = nn.Sequential(
            nn.Dropout(p = salCNN_dropoutRate)
             )

        # PrdFC Module
        prdFC_inputSize = self.gazeSeqCNN1D_outputSize + self.headSeqCNN1D_outputSize + self.objectSeqCNN1D_outputSize + self.salCNN_outputSize
        prdFC_linearSize1 = 128
        prdFC_linearSize2 = self.clusterSize
        prdFC_dropoutRate = 0.5
        # the prediction fc layer
        self.PrdFC = nn.Sequential(
            nn.Linear(prdFC_inputSize, prdFC_linearSize1),
            nn.BatchNorm1d(prdFC_linearSize1),
            nn.ReLU(),
            #nn.Sigmoid(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.Dropout(p = prdFC_dropoutRate),
            nn.Linear(prdFC_linearSize1, prdFC_linearSize2),
            nn.Softmax(dim = 1),
             )

    def forward1(self, x):
        currentGaze = x[:, 0:2]
        index = self.gazeSeqSize
        gazeSeq = x[:, 0:index]
        headSeq = x[:, index: index+self.headSeqSize]
        index += self.headSeqSize
        objectSeq = x[:, index: index+self.objectSeqSize]
        index += self.objectSeqSize
        saliencyMap = x[:, index: index+self.saliencySize]

        gazeSeq = gazeSeq.reshape(-1, self.gazeSeqLength, self.gazeFeatureNum)
        gazeSeq = gazeSeq.permute(0,2,1)
        gazeSeqOut = self.gazeSeqCNN1D(gazeSeq)
        gazeSeqOut = gazeSeqOut.reshape(-1, self.gazeSeqCNN1D_outputSize)
        gazeSeqOut = self.gazeSeqCNNDropout(gazeSeqOut)

        headSeq = headSeq.reshape(-1, self.headSeqLength, self.headFeatureNum)
        headSeq = headSeq.permute(0,2,1)
        headSeqOut = self.headSeqCNN1D(headSeq)
        headSeqOut = headSeqOut.reshape(-1, self.headSeqCNN1D_outputSize)
        headSeqOut = self.headSeqCNNDropout(headSeqOut)

        objectSeq = objectSeq.reshape(-1, self.objectSeqLength, self.objectFeatureNum)
        objectSeq = objectSeq.permute(0,2,1)
        objectSeqOut = self.objectSeqCNN1D(objectSeq)
        objectSeqOut = objectSeqOut.reshape(-1, self.objectSeqCNN1D_outputSize)
        objectSeqOut = self.objectSeqCNNDropout(objectSeqOut)

        seqOut = torch.cat((gazeSeqOut, headSeqOut), 1)
        seqOut = torch.cat((seqOut, objectSeqOut), 1)

        saliencyMap = saliencyMap.reshape(saliencyMap.size(0), self.saliencyNum, self.saliencyWidth, self.saliencyWidth)
        saliencyFeatures = self.SalCNN(saliencyMap)
        saliencyOut = saliencyFeatures.reshape(saliencyMap.size(0), -1)
        saliencyOut = self.SalCNNDropout(saliencyOut)

        prdInput = torch.cat((seqOut, saliencyOut), 1)
        prdOut = self.PrdFC(prdInput)
        out = prdOut.mm(self.cluster)
        out = out + currentGaze
        return out

    def forward(self, x):
        out = self.forward1(x)
        return out

# FixationNet for forecasting eye fixations in long-term future (over 300 ms)
# CurrentGaze is not used.
class FixationNet300(nn.Module):
    def __init__(self, gazeSeqSize, headSeqSize, objectSeqSize, saliencySize, clusterPath):
        super().__init__()
        # the input params
        self.gazeSeqSize = gazeSeqSize
        self.headSeqSize = headSeqSize
        self.objectSeqSize = objectSeqSize
        self.saliencySize = saliencySize

        print('gazeSeqSize: {}'.format(self.gazeSeqSize))
        print('headSeqSize: {}'.format(self.headSeqSize))
        print('taskObjectSeqSize: {}'.format(self.objectSeqSize))
        print('saliencySize: {}'.format(self.saliencySize))

        # preset params
        self.gazeFeatureNum = 2
        self.gazeSeqLength = int(self.gazeSeqSize/self.gazeFeatureNum)
        self.headFeatureNum = 2
        self.headSeqLength = int(self.headSeqSize/self.headFeatureNum)
        self.objectFeatureNum = 12
        self.objectSeqLength = int(self.objectSeqSize/self.objectFeatureNum)
        self.saliencyWidth = 24
        self.saliencyNum = int(self.saliencySize/(self.saliencyWidth*self.saliencyWidth))

        # load the cluster centers
        self.cluster = torch.from_numpy(np.load(clusterPath)).float().to(device)
        self.clusterSize = self.cluster.shape[0]
        print('Cluster Size: {}'.format(self.clusterSize))

        # GazeSeqCNN1D Module
        gazeSeqCNN1D_outChannels1 = 32
        gazeSeqCNN1D_poolingRate1 = 2
        gazeSeqCNN1D_kernelSize1 = 1
        gazeSeqCNN1D_featureSize1 = floor((self.gazeSeqLength - gazeSeqCNN1D_kernelSize1 + 1)/gazeSeqCNN1D_poolingRate1)
        self.gazeSeqCNN1D_outputSize = gazeSeqCNN1D_featureSize1 * gazeSeqCNN1D_outChannels1
        self.gazeSeqCNN1D = nn.Sequential(
            nn.Conv1d(in_channels=self.gazeFeatureNum, out_channels=gazeSeqCNN1D_outChannels1,kernel_size=gazeSeqCNN1D_kernelSize1),
            nn.BatchNorm1d(gazeSeqCNN1D_outChannels1),
            #nn.ReLU(),
            #nn.LeakyReLU(),
            nn.Sigmoid(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.MaxPool1d(gazeSeqCNN1D_poolingRate1),
             )

        gazeSeqCNN_dropoutRate = 0.0
        self.gazeSeqCNNDropout = nn.Sequential(
            nn.Dropout(p = gazeSeqCNN_dropoutRate)
             )

        # HeadSeqCNN1D Module
        headSeqCNN1D_outChannels1 = 64
        headSeqCNN1D_poolingRate1 = 2
        headSeqCNN1D_kernelSize1 = 1
        headSeqCNN1D_featureSize1 = floor((self.headSeqLength - headSeqCNN1D_kernelSize1 + 1)/headSeqCNN1D_poolingRate1)
        self.headSeqCNN1D_outputSize = headSeqCNN1D_featureSize1 * headSeqCNN1D_outChannels1
        self.headSeqCNN1D = nn.Sequential(
            nn.Conv1d(in_channels=self.headFeatureNum, out_channels=headSeqCNN1D_outChannels1,kernel_size=headSeqCNN1D_kernelSize1),
            nn.BatchNorm1d(headSeqCNN1D_outChannels1),
            #nn.ReLU(),
            #nn.LeakyReLU(),
            nn.Sigmoid(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.MaxPool1d(headSeqCNN1D_poolingRate1),
            #nn.Conv1d(in_channels=headSeqCNN1D_outChannels1, out_channels=headSeqCNN1D_outChannels2,kernel_size=headSeqCNN1D_kernelSize2),
            #nn.BatchNorm1d(headSeqCNN1D_outChannels2),
            #nn.Sigmoid(),
            #nn.MaxPool1d(headSeqCNN1D_poolingRate2),
             )

        headSeqCNN_dropoutRate = 0.0
        self.headSeqCNNDropout = nn.Sequential(
            nn.Dropout(p = headSeqCNN_dropoutRate)
             )

        # objectSeqCNN1D Module
        objectSeqCNN1D_outChannels1 = 64
        objectSeqCNN1D_poolingRate1 = 2
        objectSeqCNN1D_kernelSize1 = 1
        objectSeqCNN1D_featureSize1 = floor((self.objectSeqLength - objectSeqCNN1D_kernelSize1 + 1)/objectSeqCNN1D_poolingRate1)
        objectSeqCNN1D_outChannels2 = 32
        objectSeqCNN1D_poolingRate2 = 2
        objectSeqCNN1D_kernelSize2 = 1
        objectSeqCNN1D_featureSize2 = floor((objectSeqCNN1D_featureSize1 - objectSeqCNN1D_kernelSize2 + 1)/objectSeqCNN1D_poolingRate2)
        self.objectSeqCNN1D_outputSize = objectSeqCNN1D_featureSize2 * objectSeqCNN1D_outChannels2
        self.objectSeqCNN1D = nn.Sequential(
            nn.Conv1d(in_channels=self.objectFeatureNum, out_channels=objectSeqCNN1D_outChannels1,kernel_size=objectSeqCNN1D_kernelSize1),
            nn.BatchNorm1d(objectSeqCNN1D_outChannels1),
            #nn.ReLU(),
            nn.Sigmoid(),
            #nn.LeakyReLU(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.MaxPool1d(objectSeqCNN1D_poolingRate1),
            nn.Conv1d(in_channels=objectSeqCNN1D_outChannels1, out_channels=objectSeqCNN1D_outChannels2,kernel_size=objectSeqCNN1D_kernelSize2),
            nn.BatchNorm1d(objectSeqCNN1D_outChannels2),
            #nn.ReLU(),
            nn.Sigmoid(),
            nn.MaxPool1d(objectSeqCNN1D_poolingRate2),
             )

        objectSeqCNN_dropoutRate = 0.0
        self.objectSeqCNNDropout = nn.Sequential(
            nn.Dropout(p = objectSeqCNN_dropoutRate)
             )

        # SalCNN Module
        salCNN_outChannels1 = 8
        salCNN_poolingRate1 = 2
        salCNN_kernelSize1 = 1
        salCNN_padding1 = int((salCNN_kernelSize1-1)/2)
        salCNN_imageSize1 = floor((self.saliencyWidth - salCNN_kernelSize1 + 2*salCNN_padding1 + 1)/salCNN_poolingRate1)
        self.salCNN_outputSize = salCNN_imageSize1 * salCNN_imageSize1 * salCNN_outChannels1
        self.SalCNN = nn.Sequential(
            nn.Conv2d(self.saliencyNum, salCNN_outChannels1, kernel_size=salCNN_kernelSize1, padding=salCNN_padding1),
            nn.BatchNorm2d(salCNN_outChannels1),
            nn.ReLU(),
            #nn.Sigmoid(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.MaxPool2d(kernel_size=salCNN_poolingRate1),
             )

        salCNN_dropoutRate = 0.5
        self.SalCNNDropout = nn.Sequential(
            nn.Dropout(p = salCNN_dropoutRate)
             )

        # PrdFC Module
        prdFC_inputSize = self.gazeSeqCNN1D_outputSize + self.headSeqCNN1D_outputSize + self.objectSeqCNN1D_outputSize + self.salCNN_outputSize
        prdFC_linearSize1 = 128
        prdFC_linearSize2 = self.clusterSize
        prdFC_dropoutRate = 0.5
        # the prediction fc layer
        self.PrdFC = nn.Sequential(
            nn.Linear(prdFC_inputSize, prdFC_linearSize1),
            nn.BatchNorm1d(prdFC_linearSize1),
            #nn.ReLU(),
            nn.Sigmoid(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.Dropout(p = prdFC_dropoutRate),
            nn.Linear(prdFC_linearSize1, prdFC_linearSize2),
            nn.Softmax(dim = 1),
             )

    def forward1(self, x):
        #currentGaze = x[:, 0:2]
        index = self.gazeSeqSize
        gazeSeq = x[:, 0:index]
        headSeq = x[:, index: index+self.headSeqSize]
        index += self.headSeqSize
        objectSeq = x[:, index: index+self.objectSeqSize]
        index += self.objectSeqSize
        saliencyMap = x[:, index: index+self.saliencySize]

        gazeSeq = gazeSeq.reshape(-1, self.gazeSeqLength, self.gazeFeatureNum)
        gazeSeq = gazeSeq.permute(0,2,1)
        gazeSeqOut = self.gazeSeqCNN1D(gazeSeq)
        gazeSeqOut = gazeSeqOut.reshape(-1, self.gazeSeqCNN1D_outputSize)
        gazeSeqOut = self.gazeSeqCNNDropout(gazeSeqOut)

        headSeq = headSeq.reshape(-1, self.headSeqLength, self.headFeatureNum)
        headSeq = headSeq.permute(0,2,1)
        headSeqOut = self.headSeqCNN1D(headSeq)
        headSeqOut = headSeqOut.reshape(-1, self.headSeqCNN1D_outputSize)
        headSeqOut = self.headSeqCNNDropout(headSeqOut)

        objectSeq = objectSeq.reshape(-1, self.objectSeqLength, self.objectFeatureNum)
        objectSeq = objectSeq.permute(0,2,1)
        objectSeqOut = self.objectSeqCNN1D(objectSeq)
        objectSeqOut = objectSeqOut.reshape(-1, self.objectSeqCNN1D_outputSize)
        objectSeqOut = self.objectSeqCNNDropout(objectSeqOut)

        seqOut = torch.cat((gazeSeqOut, headSeqOut), 1)
        seqOut = torch.cat((seqOut, objectSeqOut), 1)

        saliencyMap = saliencyMap.reshape(saliencyMap.size(0), self.saliencyNum, self.saliencyWidth, self.saliencyWidth)
        saliencyFeatures = self.SalCNN(saliencyMap)
        saliencyOut = saliencyFeatures.reshape(saliencyMap.size(0), -1)
        saliencyOut = self.SalCNNDropout(saliencyOut)

        prdInput = torch.cat((seqOut, saliencyOut), 1)
        prdOut = self.PrdFC(prdInput)
        out = prdOut.mm(self.cluster)
        #out = out + currentGaze
        return out

    def forward(self, x):
        out = self.forward1(x)
        return out

# FixationNet for forecasting eye fixations on DGaze dataset.
# CurrentGaze is not used.
class FixationNet_DGazeDataset(nn.Module):
    def __init__(self, gazeSeqSize, headSeqSize, objectSeqSize, saliencySize, clusterPath):
        super().__init__()
        # the input params
        self.gazeSeqSize = gazeSeqSize
        self.headSeqSize = headSeqSize
        self.objectSeqSize = objectSeqSize
        self.saliencySize = saliencySize

        print('gazeSeqSize: {}'.format(self.gazeSeqSize))
        print('headSeqSize: {}'.format(self.headSeqSize))
        print('taskObjectSeqSize: {}'.format(self.objectSeqSize))
        print('saliencySize: {}'.format(self.saliencySize))

        # preset params
        self.gazeFeatureNum = 2
        self.gazeSeqLength = int(self.gazeSeqSize/self.gazeFeatureNum)
        self.headFeatureNum = 2
        self.headSeqLength = int(self.headSeqSize/self.headFeatureNum)
        # the nearest 3 dynamic objects, 3*3
        self.objectFeatureNum = 9
        self.objectSeqLength = int(self.objectSeqSize/self.objectFeatureNum)
        self.saliencyWidth = 24
        self.saliencyNum = int(self.saliencySize/(self.saliencyWidth*self.saliencyWidth))

        # load the cluster centers
        self.cluster = torch.from_numpy(np.load(clusterPath)).float().to(device)
        self.clusterSize = self.cluster.shape[0]
        print('Cluster Size: {}'.format(self.clusterSize))

        # GazeSeqCNN1D Module
        gazeSeqCNN1D_outChannels1 = 32
        gazeSeqCNN1D_poolingRate1 = 2
        gazeSeqCNN1D_kernelSize1 = 1
        gazeSeqCNN1D_featureSize1 = floor((self.gazeSeqLength - gazeSeqCNN1D_kernelSize1 + 1)/gazeSeqCNN1D_poolingRate1)
        self.gazeSeqCNN1D_outputSize = gazeSeqCNN1D_featureSize1 * gazeSeqCNN1D_outChannels1
        self.gazeSeqCNN1D = nn.Sequential(
            nn.Conv1d(in_channels=self.gazeFeatureNum, out_channels=gazeSeqCNN1D_outChannels1,kernel_size=gazeSeqCNN1D_kernelSize1),
            nn.BatchNorm1d(gazeSeqCNN1D_outChannels1),
            #nn.ReLU(),
            #nn.LeakyReLU(),
            nn.Sigmoid(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.MaxPool1d(gazeSeqCNN1D_poolingRate1),
             )

        gazeSeqCNN_dropoutRate = 0.0
        self.gazeSeqCNNDropout = nn.Sequential(
            nn.Dropout(p = gazeSeqCNN_dropoutRate)
             )

        # HeadSeqCNN1D Module
        headSeqCNN1D_outChannels1 = 64
        headSeqCNN1D_poolingRate1 = 2
        headSeqCNN1D_kernelSize1 = 1
        headSeqCNN1D_featureSize1 = floor((self.headSeqLength - headSeqCNN1D_kernelSize1 + 1)/headSeqCNN1D_poolingRate1)
        self.headSeqCNN1D_outputSize = headSeqCNN1D_featureSize1 * headSeqCNN1D_outChannels1
        self.headSeqCNN1D = nn.Sequential(
            nn.Conv1d(in_channels=self.headFeatureNum, out_channels=headSeqCNN1D_outChannels1,kernel_size=headSeqCNN1D_kernelSize1),
            nn.BatchNorm1d(headSeqCNN1D_outChannels1),
            #nn.ReLU(),
            #nn.LeakyReLU(),
            nn.Sigmoid(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.MaxPool1d(headSeqCNN1D_poolingRate1),
            #nn.Conv1d(in_channels=headSeqCNN1D_outChannels1, out_channels=headSeqCNN1D_outChannels2,kernel_size=headSeqCNN1D_kernelSize2),
            #nn.BatchNorm1d(headSeqCNN1D_outChannels2),
            #nn.Sigmoid(),
            #nn.MaxPool1d(headSeqCNN1D_poolingRate2),
             )

        headSeqCNN_dropoutRate = 0.0
        self.headSeqCNNDropout = nn.Sequential(
            nn.Dropout(p = headSeqCNN_dropoutRate)
             )

        # objectSeqCNN1D Module
        objectSeqCNN1D_outChannels1 = 64
        objectSeqCNN1D_poolingRate1 = 2
        objectSeqCNN1D_kernelSize1 = 1
        objectSeqCNN1D_featureSize1 = floor((self.objectSeqLength - objectSeqCNN1D_kernelSize1 + 1)/objectSeqCNN1D_poolingRate1)
        objectSeqCNN1D_outChannels2 = 32
        objectSeqCNN1D_poolingRate2 = 2
        objectSeqCNN1D_kernelSize2 = 1
        objectSeqCNN1D_featureSize2 = floor((objectSeqCNN1D_featureSize1 - objectSeqCNN1D_kernelSize2 + 1)/objectSeqCNN1D_poolingRate2)
        self.objectSeqCNN1D_outputSize = objectSeqCNN1D_featureSize2 * objectSeqCNN1D_outChannels2
        self.objectSeqCNN1D = nn.Sequential(
            nn.Conv1d(in_channels=self.objectFeatureNum, out_channels=objectSeqCNN1D_outChannels1,kernel_size=objectSeqCNN1D_kernelSize1),
            nn.BatchNorm1d(objectSeqCNN1D_outChannels1),
            #nn.ReLU(),
            nn.Sigmoid(),
            #nn.LeakyReLU(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.MaxPool1d(objectSeqCNN1D_poolingRate1),
            nn.Conv1d(in_channels=objectSeqCNN1D_outChannels1, out_channels=objectSeqCNN1D_outChannels2,kernel_size=objectSeqCNN1D_kernelSize2),
            nn.BatchNorm1d(objectSeqCNN1D_outChannels2),
            #nn.ReLU(),
            nn.Sigmoid(),
            nn.MaxPool1d(objectSeqCNN1D_poolingRate2),
             )

        objectSeqCNN_dropoutRate = 0.0
        self.objectSeqCNNDropout = nn.Sequential(
            nn.Dropout(p = objectSeqCNN_dropoutRate)
             )

        # SalCNN Module
        salCNN_outChannels1 = 8
        salCNN_poolingRate1 = 2
        salCNN_kernelSize1 = 1
        salCNN_padding1 = int((salCNN_kernelSize1-1)/2)
        salCNN_imageSize1 = floor((self.saliencyWidth - salCNN_kernelSize1 + 2*salCNN_padding1 + 1)/salCNN_poolingRate1)
        self.salCNN_outputSize = salCNN_imageSize1 * salCNN_imageSize1 * salCNN_outChannels1
        self.SalCNN = nn.Sequential(
            nn.Conv2d(self.saliencyNum, salCNN_outChannels1, kernel_size=salCNN_kernelSize1, padding=salCNN_padding1),
            nn.BatchNorm2d(salCNN_outChannels1),
            nn.ReLU(),
            #nn.Sigmoid(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.MaxPool2d(kernel_size=salCNN_poolingRate1),
             )

        salCNN_dropoutRate = 0.5
        self.SalCNNDropout = nn.Sequential(
            nn.Dropout(p = salCNN_dropoutRate)
             )

        # PrdFC Module
        prdFC_inputSize = self.gazeSeqCNN1D_outputSize + self.headSeqCNN1D_outputSize + self.objectSeqCNN1D_outputSize + self.salCNN_outputSize
        prdFC_linearSize1 = 128
        prdFC_linearSize2 = self.clusterSize
        prdFC_dropoutRate = 0.5
        # the prediction fc layer
        self.PrdFC = nn.Sequential(
            nn.Linear(prdFC_inputSize, prdFC_linearSize1),
            nn.BatchNorm1d(prdFC_linearSize1),
            #nn.ReLU(),
            nn.Sigmoid(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.Dropout(p = prdFC_dropoutRate),
            nn.Linear(prdFC_linearSize1, prdFC_linearSize2),
            nn.Softmax(dim = 1),
             )

    def forward1(self, x):
        #currentGaze = x[:, 0:2]
        index = self.gazeSeqSize
        gazeSeq = x[:, 0:index]
        headSeq = x[:, index: index+self.headSeqSize]
        index += self.headSeqSize
        objectSeq = x[:, index: index+self.objectSeqSize]
        index += self.objectSeqSize
        saliencyMap = x[:, index: index+self.saliencySize]

        gazeSeq = gazeSeq.reshape(-1, self.gazeSeqLength, self.gazeFeatureNum)
        gazeSeq = gazeSeq.permute(0,2,1)
        gazeSeqOut = self.gazeSeqCNN1D(gazeSeq)
        gazeSeqOut = gazeSeqOut.reshape(-1, self.gazeSeqCNN1D_outputSize)
        gazeSeqOut = self.gazeSeqCNNDropout(gazeSeqOut)

        headSeq = headSeq.reshape(-1, self.headSeqLength, self.headFeatureNum)
        headSeq = headSeq.permute(0,2,1)
        headSeqOut = self.headSeqCNN1D(headSeq)
        headSeqOut = headSeqOut.reshape(-1, self.headSeqCNN1D_outputSize)
        headSeqOut = self.headSeqCNNDropout(headSeqOut)

        objectSeq = objectSeq.reshape(-1, self.objectSeqLength, self.objectFeatureNum)
        objectSeq = objectSeq.permute(0,2,1)
        objectSeqOut = self.objectSeqCNN1D(objectSeq)
        objectSeqOut = objectSeqOut.reshape(-1, self.objectSeqCNN1D_outputSize)
        objectSeqOut = self.objectSeqCNNDropout(objectSeqOut)

        seqOut = torch.cat((gazeSeqOut, headSeqOut), 1)
        seqOut = torch.cat((seqOut, objectSeqOut), 1)

        saliencyMap = saliencyMap.reshape(saliencyMap.size(0), self.saliencyNum, self.saliencyWidth, self.saliencyWidth)
        saliencyFeatures = self.SalCNN(saliencyMap)
        saliencyOut = saliencyFeatures.reshape(saliencyMap.size(0), -1)
        saliencyOut = self.SalCNNDropout(saliencyOut)

        prdInput = torch.cat((seqOut, saliencyOut), 1)
        prdOut = self.PrdFC(prdInput)
        out = prdOut.mm(self.cluster)
        #out = out + currentGaze
        return out

    def forward(self, x):
        out = self.forward1(x)
        return out

class FixationNet_without_Saliency(nn.Module):
    def __init__(self, gazeSeqSize, headSeqSize, objectSeqSize, saliencySize, clusterPath):
        super().__init__()
        # the input params
        self.gazeSeqSize = gazeSeqSize
        self.headSeqSize = headSeqSize
        self.objectSeqSize = objectSeqSize
        self.saliencySize = saliencySize

        print('gazeSeqSize: {}'.format(self.gazeSeqSize))
        print('headSeqSize: {}'.format(self.headSeqSize))
        print('taskObjectSeqSize: {}'.format(self.objectSeqSize))
        print('saliencySize: {}'.format(self.saliencySize))

        # preset params
        self.gazeFeatureNum = 2
        self.gazeSeqLength = int(self.gazeSeqSize/self.gazeFeatureNum)
        self.headFeatureNum = 2
        self.headSeqLength = int(self.headSeqSize/self.headFeatureNum)
        self.objectFeatureNum = 12
        self.objectSeqLength = int(self.objectSeqSize/self.objectFeatureNum)
        self.saliencyWidth = 24
        self.saliencyNum = int(self.saliencySize/(self.saliencyWidth*self.saliencyWidth))

        # load the cluster centers
        self.cluster = torch.from_numpy(np.load(clusterPath)).float().to(device)
        self.clusterSize = self.cluster.shape[0]
        print('Cluster Size: {}'.format(self.clusterSize))

        # GazeSeqCNN1D Module
        gazeSeqCNN1D_outChannels1 = 32
        gazeSeqCNN1D_poolingRate1 = 2
        gazeSeqCNN1D_kernelSize1 = 1
        gazeSeqCNN1D_featureSize1 = floor((self.gazeSeqLength - gazeSeqCNN1D_kernelSize1 + 1)/gazeSeqCNN1D_poolingRate1)
        self.gazeSeqCNN1D_outputSize = gazeSeqCNN1D_featureSize1 * gazeSeqCNN1D_outChannels1
        self.gazeSeqCNN1D = nn.Sequential(
            nn.Conv1d(in_channels=self.gazeFeatureNum, out_channels=gazeSeqCNN1D_outChannels1,kernel_size=gazeSeqCNN1D_kernelSize1),
            nn.BatchNorm1d(gazeSeqCNN1D_outChannels1),
            nn.ReLU(),
            #nn.LeakyReLU(),
            #nn.Sigmoid(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.MaxPool1d(gazeSeqCNN1D_poolingRate1),
             )

        gazeSeqCNN_dropoutRate = 0.0
        self.gazeSeqCNNDropout = nn.Sequential(
            nn.Dropout(p = gazeSeqCNN_dropoutRate)
             )

        # HeadSeqCNN1D Module
        headSeqCNN1D_outChannels1 = 64
        headSeqCNN1D_poolingRate1 = 2
        headSeqCNN1D_kernelSize1 = 1
        headSeqCNN1D_featureSize1 = floor((self.headSeqLength - headSeqCNN1D_kernelSize1 + 1)/headSeqCNN1D_poolingRate1)
        self.headSeqCNN1D_outputSize = headSeqCNN1D_featureSize1 * headSeqCNN1D_outChannels1
        self.headSeqCNN1D = nn.Sequential(
            nn.Conv1d(in_channels=self.headFeatureNum, out_channels=headSeqCNN1D_outChannels1,kernel_size=headSeqCNN1D_kernelSize1),
            nn.BatchNorm1d(headSeqCNN1D_outChannels1),
            nn.ReLU(),
            #nn.LeakyReLU(),
            #nn.Sigmoid(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.MaxPool1d(headSeqCNN1D_poolingRate1),
            #nn.Conv1d(in_channels=headSeqCNN1D_outChannels1, out_channels=headSeqCNN1D_outChannels2,kernel_size=headSeqCNN1D_kernelSize2),
            #nn.BatchNorm1d(headSeqCNN1D_outChannels2),
            #nn.Sigmoid(),
            #nn.MaxPool1d(headSeqCNN1D_poolingRate2),
             )

        headSeqCNN_dropoutRate = 0.0
        self.headSeqCNNDropout = nn.Sequential(
            nn.Dropout(p = headSeqCNN_dropoutRate)
             )

        # objectSeqCNN1D Module
        objectSeqCNN1D_outChannels1 = 64
        objectSeqCNN1D_poolingRate1 = 2
        objectSeqCNN1D_kernelSize1 = 1
        objectSeqCNN1D_featureSize1 = floor((self.objectSeqLength - objectSeqCNN1D_kernelSize1 + 1)/objectSeqCNN1D_poolingRate1)
        objectSeqCNN1D_outChannels2 = 32
        objectSeqCNN1D_poolingRate2 = 2
        objectSeqCNN1D_kernelSize2 = 1
        objectSeqCNN1D_featureSize2 = floor((objectSeqCNN1D_featureSize1 - objectSeqCNN1D_kernelSize2 + 1)/objectSeqCNN1D_poolingRate2)
        self.objectSeqCNN1D_outputSize = objectSeqCNN1D_featureSize2 * objectSeqCNN1D_outChannels2
        self.objectSeqCNN1D = nn.Sequential(
            nn.Conv1d(in_channels=self.objectFeatureNum, out_channels=objectSeqCNN1D_outChannels1,kernel_size=objectSeqCNN1D_kernelSize1),
            nn.BatchNorm1d(objectSeqCNN1D_outChannels1),
            nn.ReLU(),
            #nn.Sigmoid(),
            #nn.LeakyReLU(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.MaxPool1d(objectSeqCNN1D_poolingRate1),
            nn.Conv1d(in_channels=objectSeqCNN1D_outChannels1, out_channels=objectSeqCNN1D_outChannels2,kernel_size=objectSeqCNN1D_kernelSize2),
            nn.BatchNorm1d(objectSeqCNN1D_outChannels2),
            nn.ReLU(),
            #nn.Sigmoid(),
            nn.MaxPool1d(objectSeqCNN1D_poolingRate2),
             )

        objectSeqCNN_dropoutRate = 0.0
        self.objectSeqCNNDropout = nn.Sequential(
            nn.Dropout(p = objectSeqCNN_dropoutRate)
             )

        # SalCNN Module
        salCNN_outChannels1 = 8
        salCNN_poolingRate1 = 2
        salCNN_kernelSize1 = 1
        salCNN_padding1 = int((salCNN_kernelSize1-1)/2)
        salCNN_imageSize1 = floor((self.saliencyWidth - salCNN_kernelSize1 + 2*salCNN_padding1 + 1)/salCNN_poolingRate1)
        self.salCNN_outputSize = salCNN_imageSize1 * salCNN_imageSize1 * salCNN_outChannels1
        self.SalCNN = nn.Sequential(
            nn.Conv2d(self.saliencyNum, salCNN_outChannels1, kernel_size=salCNN_kernelSize1, padding=salCNN_padding1),
            nn.BatchNorm2d(salCNN_outChannels1),
            nn.ReLU(),
            #nn.Sigmoid(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.MaxPool2d(kernel_size=salCNN_poolingRate1),
             )

        salCNN_dropoutRate = 0.5
        self.SalCNNDropout = nn.Sequential(
            nn.Dropout(p = salCNN_dropoutRate)
             )

        # PrdFC Module
        prdFC_inputSize = self.gazeSeqCNN1D_outputSize + self.headSeqCNN1D_outputSize + self.objectSeqCNN1D_outputSize
        prdFC_linearSize1 = 128
        #prdFC_linearSize2 = 128
        prdFC_linearSize2 = self.clusterSize
        prdFC_dropoutRate = 0.5
        # the prediction fc layer for DGaze.
        self.PrdFC = nn.Sequential(
            nn.Linear(prdFC_inputSize, prdFC_linearSize1),
            nn.BatchNorm1d(prdFC_linearSize1),
            nn.ReLU(),
            #nn.Sigmoid(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.Dropout(p = prdFC_dropoutRate),
            nn.Linear(prdFC_linearSize1, prdFC_linearSize2),
            #nn.BatchNorm1d(prdFC_linearSize2),
            #nn.Sigmoid(),
            #nn.Dropout(p = prdFC_dropoutRate),
            #nn.Linear(prdFC_linearSize2, 2),
            nn.Softmax(dim = 1),
             )

    def forward1(self, x):
        currentGaze = x[:, 0:2]
        index = self.gazeSeqSize
        gazeSeq = x[:, 0:index]
        headSeq = x[:, index: index+self.headSeqSize]
        index += self.headSeqSize
        objectSeq = x[:, index: index+self.objectSeqSize]
        index += self.objectSeqSize
        #saliencyMap = x[:, index: index+self.saliencySize]

        gazeSeq = gazeSeq.reshape(-1, self.gazeSeqLength, self.gazeFeatureNum)
        gazeSeq = gazeSeq.permute(0,2,1)
        gazeSeqOut = self.gazeSeqCNN1D(gazeSeq)
        gazeSeqOut = gazeSeqOut.reshape(-1, self.gazeSeqCNN1D_outputSize)
        gazeSeqOut = self.gazeSeqCNNDropout(gazeSeqOut)

        headSeq = headSeq.reshape(-1, self.headSeqLength, self.headFeatureNum)
        headSeq = headSeq.permute(0,2,1)
        headSeqOut = self.headSeqCNN1D(headSeq)
        headSeqOut = headSeqOut.reshape(-1, self.headSeqCNN1D_outputSize)
        headSeqOut = self.headSeqCNNDropout(headSeqOut)

        objectSeq = objectSeq.reshape(-1, self.objectSeqLength, self.objectFeatureNum)
        objectSeq = objectSeq.permute(0,2,1)
        objectSeqOut = self.objectSeqCNN1D(objectSeq)
        objectSeqOut = objectSeqOut.reshape(-1, self.objectSeqCNN1D_outputSize)
        objectSeqOut = self.objectSeqCNNDropout(objectSeqOut)

        seqOut = torch.cat((gazeSeqOut, headSeqOut), 1)
        seqOut = torch.cat((seqOut, objectSeqOut), 1)

        #saliencyMap = saliencyMap.reshape(saliencyMap.size(0), self.saliencyNum, self.saliencyWidth, self.saliencyWidth)
        #saliencyFeatures = self.SalCNN(saliencyMap)
        #saliencyOut = saliencyFeatures.reshape(saliencyMap.size(0), -1)
        #saliencyOut = self.SalCNNDropout(saliencyOut)

        prdInput = seqOut
        prdOut = self.PrdFC(prdInput)
        out = prdOut.mm(self.cluster)
        out = out + currentGaze
        return out

    def forward(self, x):
        out = self.forward1(x)
        return out

class FixationNet_without_Task(nn.Module):
    def __init__(self, gazeSeqSize, headSeqSize, objectSeqSize, saliencySize, clusterPath):
        super().__init__()
        # the input params
        self.gazeSeqSize = gazeSeqSize
        self.headSeqSize = headSeqSize
        self.objectSeqSize = objectSeqSize
        self.saliencySize = saliencySize

        print('gazeSeqSize: {}'.format(self.gazeSeqSize))
        print('headSeqSize: {}'.format(self.headSeqSize))
        print('taskObjectSeqSize: {}'.format(self.objectSeqSize))
        print('saliencySize: {}'.format(self.saliencySize))

        # preset params
        self.gazeFeatureNum = 2
        self.gazeSeqLength = int(self.gazeSeqSize/self.gazeFeatureNum)
        self.headFeatureNum = 2
        self.headSeqLength = int(self.headSeqSize/self.headFeatureNum)
        self.objectFeatureNum = 12
        self.objectSeqLength = int(self.objectSeqSize/self.objectFeatureNum)
        self.saliencyWidth = 24
        self.saliencyNum = int(self.saliencySize/(self.saliencyWidth*self.saliencyWidth))

        # load the cluster centers
        self.cluster = torch.from_numpy(np.load(clusterPath)).float().to(device)
        self.clusterSize = self.cluster.shape[0]
        print('Cluster Size: {}'.format(self.clusterSize))

        # GazeSeqCNN1D Module
        gazeSeqCNN1D_outChannels1 = 32
        gazeSeqCNN1D_poolingRate1 = 2
        gazeSeqCNN1D_kernelSize1 = 1
        gazeSeqCNN1D_featureSize1 = floor((self.gazeSeqLength - gazeSeqCNN1D_kernelSize1 + 1)/gazeSeqCNN1D_poolingRate1)
        self.gazeSeqCNN1D_outputSize = gazeSeqCNN1D_featureSize1 * gazeSeqCNN1D_outChannels1
        self.gazeSeqCNN1D = nn.Sequential(
            nn.Conv1d(in_channels=self.gazeFeatureNum, out_channels=gazeSeqCNN1D_outChannels1,kernel_size=gazeSeqCNN1D_kernelSize1),
            nn.BatchNorm1d(gazeSeqCNN1D_outChannels1),
            nn.ReLU(),
            #nn.LeakyReLU(),
            #nn.Sigmoid(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.MaxPool1d(gazeSeqCNN1D_poolingRate1),
             )

        gazeSeqCNN_dropoutRate = 0.0
        self.gazeSeqCNNDropout = nn.Sequential(
            nn.Dropout(p = gazeSeqCNN_dropoutRate)
             )

        # HeadSeqCNN1D Module
        headSeqCNN1D_outChannels1 = 64
        headSeqCNN1D_poolingRate1 = 2
        headSeqCNN1D_kernelSize1 = 1
        headSeqCNN1D_featureSize1 = floor((self.headSeqLength - headSeqCNN1D_kernelSize1 + 1)/headSeqCNN1D_poolingRate1)
        self.headSeqCNN1D_outputSize = headSeqCNN1D_featureSize1 * headSeqCNN1D_outChannels1
        self.headSeqCNN1D = nn.Sequential(
            nn.Conv1d(in_channels=self.headFeatureNum, out_channels=headSeqCNN1D_outChannels1,kernel_size=headSeqCNN1D_kernelSize1),
            nn.BatchNorm1d(headSeqCNN1D_outChannels1),
            nn.ReLU(),
            #nn.LeakyReLU(),
            #nn.Sigmoid(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.MaxPool1d(headSeqCNN1D_poolingRate1),
            #nn.Conv1d(in_channels=headSeqCNN1D_outChannels1, out_channels=headSeqCNN1D_outChannels2,kernel_size=headSeqCNN1D_kernelSize2),
            #nn.BatchNorm1d(headSeqCNN1D_outChannels2),
            #nn.Sigmoid(),
            #nn.MaxPool1d(headSeqCNN1D_poolingRate2),
             )

        headSeqCNN_dropoutRate = 0.0
        self.headSeqCNNDropout = nn.Sequential(
            nn.Dropout(p = headSeqCNN_dropoutRate)
             )

        # objectSeqCNN1D Module
        objectSeqCNN1D_outChannels1 = 64
        objectSeqCNN1D_poolingRate1 = 2
        objectSeqCNN1D_kernelSize1 = 1
        objectSeqCNN1D_featureSize1 = floor((self.objectSeqLength - objectSeqCNN1D_kernelSize1 + 1)/objectSeqCNN1D_poolingRate1)
        objectSeqCNN1D_outChannels2 = 32
        objectSeqCNN1D_poolingRate2 = 2
        objectSeqCNN1D_kernelSize2 = 1
        objectSeqCNN1D_featureSize2 = floor((objectSeqCNN1D_featureSize1 - objectSeqCNN1D_kernelSize2 + 1)/objectSeqCNN1D_poolingRate2)
        self.objectSeqCNN1D_outputSize = objectSeqCNN1D_featureSize2 * objectSeqCNN1D_outChannels2
        self.objectSeqCNN1D = nn.Sequential(
            nn.Conv1d(in_channels=self.objectFeatureNum, out_channels=objectSeqCNN1D_outChannels1,kernel_size=objectSeqCNN1D_kernelSize1),
            nn.BatchNorm1d(objectSeqCNN1D_outChannels1),
            nn.ReLU(),
            #nn.Sigmoid(),
            #nn.LeakyReLU(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.MaxPool1d(objectSeqCNN1D_poolingRate1),
            nn.Conv1d(in_channels=objectSeqCNN1D_outChannels1, out_channels=objectSeqCNN1D_outChannels2,kernel_size=objectSeqCNN1D_kernelSize2),
            nn.BatchNorm1d(objectSeqCNN1D_outChannels2),
            nn.ReLU(),
            #nn.Sigmoid(),
            nn.MaxPool1d(objectSeqCNN1D_poolingRate2),
             )

        objectSeqCNN_dropoutRate = 0.0
        self.objectSeqCNNDropout = nn.Sequential(
            nn.Dropout(p = objectSeqCNN_dropoutRate)
             )

        # SalCNN Module
        salCNN_outChannels1 = 8
        salCNN_poolingRate1 = 2
        salCNN_kernelSize1 = 1
        salCNN_padding1 = int((salCNN_kernelSize1-1)/2)
        salCNN_imageSize1 = floor((self.saliencyWidth - salCNN_kernelSize1 + 2*salCNN_padding1 + 1)/salCNN_poolingRate1)
        self.salCNN_outputSize = salCNN_imageSize1 * salCNN_imageSize1 * salCNN_outChannels1
        self.SalCNN = nn.Sequential(
            nn.Conv2d(self.saliencyNum, salCNN_outChannels1, kernel_size=salCNN_kernelSize1, padding=salCNN_padding1),
            nn.BatchNorm2d(salCNN_outChannels1),
            nn.ReLU(),
            #nn.Sigmoid(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.MaxPool2d(kernel_size=salCNN_poolingRate1),
             )

        salCNN_dropoutRate = 0.5
        self.SalCNNDropout = nn.Sequential(
            nn.Dropout(p = salCNN_dropoutRate)
             )

        # PrdFC Module
        prdFC_inputSize = self.gazeSeqCNN1D_outputSize + self.headSeqCNN1D_outputSize + self.salCNN_outputSize
        prdFC_linearSize1 = 128
        prdFC_linearSize2 = self.clusterSize
        prdFC_dropoutRate = 0.5
        # the prediction fc layer for DGaze.
        self.PrdFC = nn.Sequential(
            nn.Linear(prdFC_inputSize, prdFC_linearSize1),
            nn.BatchNorm1d(prdFC_linearSize1),
            nn.ReLU(),
            #nn.Sigmoid(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.Dropout(p = prdFC_dropoutRate),
            nn.Linear(prdFC_linearSize1, prdFC_linearSize2),
            nn.Softmax(dim = 1),
             )

    def forward1(self, x):
        currentGaze = x[:, 0:2]
        index = self.gazeSeqSize
        gazeSeq = x[:, 0:index]
        headSeq = x[:, index: index+self.headSeqSize]
        index += self.headSeqSize
        #objectSeq = x[:, index: index+self.objectSeqSize]
        index += self.objectSeqSize
        saliencyMap = x[:, index: index+self.saliencySize]

        gazeSeq = gazeSeq.reshape(-1, self.gazeSeqLength, self.gazeFeatureNum)
        gazeSeq = gazeSeq.permute(0,2,1)
        gazeSeqOut = self.gazeSeqCNN1D(gazeSeq)
        gazeSeqOut = gazeSeqOut.reshape(-1, self.gazeSeqCNN1D_outputSize)
        gazeSeqOut = self.gazeSeqCNNDropout(gazeSeqOut)

        headSeq = headSeq.reshape(-1, self.headSeqLength, self.headFeatureNum)
        headSeq = headSeq.permute(0,2,1)
        headSeqOut = self.headSeqCNN1D(headSeq)
        headSeqOut = headSeqOut.reshape(-1, self.headSeqCNN1D_outputSize)
        headSeqOut = self.headSeqCNNDropout(headSeqOut)

        #objectSeq = objectSeq.reshape(-1, self.objectSeqLength, self.objectFeatureNum)
        #objectSeq = objectSeq.permute(0,2,1)
        #objectSeqOut = self.objectSeqCNN1D(objectSeq)
        #objectSeqOut = objectSeqOut.reshape(-1, self.objectSeqCNN1D_outputSize)
        #objectSeqOut = self.objectSeqCNNDropout(objectSeqOut)

        seqOut = torch.cat((gazeSeqOut, headSeqOut), 1)

        saliencyMap = saliencyMap.reshape(saliencyMap.size(0), self.saliencyNum, self.saliencyWidth, self.saliencyWidth)
        saliencyFeatures = self.SalCNN(saliencyMap)
        saliencyOut = saliencyFeatures.reshape(saliencyMap.size(0), -1)
        saliencyOut = self.SalCNNDropout(saliencyOut)

        prdInput = torch.cat((seqOut, saliencyOut), 1)
        prdOut = self.PrdFC(prdInput)
        out = prdOut.mm(self.cluster)
        out = out + currentGaze
        return out

    def forward(self, x):
        out = self.forward1(x)
        return out

class FixationNet_without_Gaze(nn.Module):
    def __init__(self, gazeSeqSize, headSeqSize, objectSeqSize, saliencySize, clusterPath):
        super().__init__()
        # the input params
        self.gazeSeqSize = gazeSeqSize
        self.headSeqSize = headSeqSize
        self.objectSeqSize = objectSeqSize
        self.saliencySize = saliencySize

        print('gazeSeqSize: {}'.format(self.gazeSeqSize))
        print('headSeqSize: {}'.format(self.headSeqSize))
        print('taskObjectSeqSize: {}'.format(self.objectSeqSize))
        print('saliencySize: {}'.format(self.saliencySize))

        # preset params
        self.gazeFeatureNum = 2
        self.gazeSeqLength = int(self.gazeSeqSize/self.gazeFeatureNum)
        self.headFeatureNum = 2
        self.headSeqLength = int(self.headSeqSize/self.headFeatureNum)
        self.objectFeatureNum = 12
        self.objectSeqLength = int(self.objectSeqSize/self.objectFeatureNum)
        self.saliencyWidth = 24
        self.saliencyNum = int(self.saliencySize/(self.saliencyWidth*self.saliencyWidth))

        # load the cluster centers
        self.cluster = torch.from_numpy(np.load(clusterPath)).float().to(device)
        self.clusterSize = self.cluster.shape[0]
        print('Cluster Size: {}'.format(self.clusterSize))

        # GazeSeqCNN1D Module
        gazeSeqCNN1D_outChannels1 = 32
        gazeSeqCNN1D_poolingRate1 = 2
        gazeSeqCNN1D_kernelSize1 = 1
        gazeSeqCNN1D_featureSize1 = floor((self.gazeSeqLength - gazeSeqCNN1D_kernelSize1 + 1)/gazeSeqCNN1D_poolingRate1)
        self.gazeSeqCNN1D_outputSize = gazeSeqCNN1D_featureSize1 * gazeSeqCNN1D_outChannels1
        self.gazeSeqCNN1D = nn.Sequential(
            nn.Conv1d(in_channels=self.gazeFeatureNum, out_channels=gazeSeqCNN1D_outChannels1,kernel_size=gazeSeqCNN1D_kernelSize1),
            nn.BatchNorm1d(gazeSeqCNN1D_outChannels1),
            nn.ReLU(),
            #nn.LeakyReLU(),
            #nn.Sigmoid(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.MaxPool1d(gazeSeqCNN1D_poolingRate1),
             )

        gazeSeqCNN_dropoutRate = 0.0
        self.gazeSeqCNNDropout = nn.Sequential(
            nn.Dropout(p = gazeSeqCNN_dropoutRate)
             )

        # HeadSeqCNN1D Module
        headSeqCNN1D_outChannels1 = 64
        headSeqCNN1D_poolingRate1 = 2
        headSeqCNN1D_kernelSize1 = 1
        headSeqCNN1D_featureSize1 = floor((self.headSeqLength - headSeqCNN1D_kernelSize1 + 1)/headSeqCNN1D_poolingRate1)
        self.headSeqCNN1D_outputSize = headSeqCNN1D_featureSize1 * headSeqCNN1D_outChannels1
        self.headSeqCNN1D = nn.Sequential(
            nn.Conv1d(in_channels=self.headFeatureNum, out_channels=headSeqCNN1D_outChannels1,kernel_size=headSeqCNN1D_kernelSize1),
            nn.BatchNorm1d(headSeqCNN1D_outChannels1),
            nn.ReLU(),
            #nn.LeakyReLU(),
            #nn.Sigmoid(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.MaxPool1d(headSeqCNN1D_poolingRate1),
            #nn.Conv1d(in_channels=headSeqCNN1D_outChannels1, out_channels=headSeqCNN1D_outChannels2,kernel_size=headSeqCNN1D_kernelSize2),
            #nn.BatchNorm1d(headSeqCNN1D_outChannels2),
            #nn.Sigmoid(),
            #nn.MaxPool1d(headSeqCNN1D_poolingRate2),
             )

        headSeqCNN_dropoutRate = 0.0
        self.headSeqCNNDropout = nn.Sequential(
            nn.Dropout(p = headSeqCNN_dropoutRate)
             )

        # objectSeqCNN1D Module
        objectSeqCNN1D_outChannels1 = 64
        objectSeqCNN1D_poolingRate1 = 2
        objectSeqCNN1D_kernelSize1 = 1
        objectSeqCNN1D_featureSize1 = floor((self.objectSeqLength - objectSeqCNN1D_kernelSize1 + 1)/objectSeqCNN1D_poolingRate1)
        objectSeqCNN1D_outChannels2 = 32
        objectSeqCNN1D_poolingRate2 = 2
        objectSeqCNN1D_kernelSize2 = 1
        objectSeqCNN1D_featureSize2 = floor((objectSeqCNN1D_featureSize1 - objectSeqCNN1D_kernelSize2 + 1)/objectSeqCNN1D_poolingRate2)
        self.objectSeqCNN1D_outputSize = objectSeqCNN1D_featureSize2 * objectSeqCNN1D_outChannels2
        self.objectSeqCNN1D = nn.Sequential(
            nn.Conv1d(in_channels=self.objectFeatureNum, out_channels=objectSeqCNN1D_outChannels1,kernel_size=objectSeqCNN1D_kernelSize1),
            nn.BatchNorm1d(objectSeqCNN1D_outChannels1),
            nn.ReLU(),
            #nn.Sigmoid(),
            #nn.LeakyReLU(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.MaxPool1d(objectSeqCNN1D_poolingRate1),
            nn.Conv1d(in_channels=objectSeqCNN1D_outChannels1, out_channels=objectSeqCNN1D_outChannels2,kernel_size=objectSeqCNN1D_kernelSize2),
            nn.BatchNorm1d(objectSeqCNN1D_outChannels2),
            nn.ReLU(),
            #nn.Sigmoid(),
            nn.MaxPool1d(objectSeqCNN1D_poolingRate2),
             )

        objectSeqCNN_dropoutRate = 0.0
        self.objectSeqCNNDropout = nn.Sequential(
            nn.Dropout(p = objectSeqCNN_dropoutRate)
             )

        # SalCNN Module
        salCNN_outChannels1 = 8
        salCNN_poolingRate1 = 2
        salCNN_kernelSize1 = 1
        salCNN_padding1 = int((salCNN_kernelSize1-1)/2)
        salCNN_imageSize1 = floor((self.saliencyWidth - salCNN_kernelSize1 + 2*salCNN_padding1 + 1)/salCNN_poolingRate1)
        self.salCNN_outputSize = salCNN_imageSize1 * salCNN_imageSize1 * salCNN_outChannels1
        self.SalCNN = nn.Sequential(
            nn.Conv2d(self.saliencyNum, salCNN_outChannels1, kernel_size=salCNN_kernelSize1, padding=salCNN_padding1),
            nn.BatchNorm2d(salCNN_outChannels1),
            nn.ReLU(),
            #nn.Sigmoid(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.MaxPool2d(kernel_size=salCNN_poolingRate1),
             )

        salCNN_dropoutRate = 0.5
        self.SalCNNDropout = nn.Sequential(
            nn.Dropout(p = salCNN_dropoutRate)
             )

        # PrdFC Module
        prdFC_inputSize = self.headSeqCNN1D_outputSize + self.objectSeqCNN1D_outputSize + self.salCNN_outputSize
        prdFC_linearSize1 = 128
        prdFC_linearSize2 = self.clusterSize
        prdFC_dropoutRate = 0.5
        # the prediction fc layer for DGaze.
        self.PrdFC = nn.Sequential(
            nn.Linear(prdFC_inputSize, prdFC_linearSize1),
            nn.BatchNorm1d(prdFC_linearSize1),
            nn.ReLU(),
            #nn.Sigmoid(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.Dropout(p = prdFC_dropoutRate),
            nn.Linear(prdFC_linearSize1, prdFC_linearSize2),
            nn.Softmax(dim = 1),
             )

    def forward1(self, x):
        currentGaze = x[:, 0:2]
        index = self.gazeSeqSize
        #gazeSeq = x[:, 0:index]
        headSeq = x[:, index: index+self.headSeqSize]
        index += self.headSeqSize
        objectSeq = x[:, index: index+self.objectSeqSize]
        index += self.objectSeqSize
        saliencyMap = x[:, index: index+self.saliencySize]

        #gazeSeq = gazeSeq.reshape(-1, self.gazeSeqLength, self.gazeFeatureNum)
        #gazeSeq = gazeSeq.permute(0,2,1)
        #gazeSeqOut = self.gazeSeqCNN1D(gazeSeq)
        #gazeSeqOut = gazeSeqOut.reshape(-1, self.gazeSeqCNN1D_outputSize)
        #gazeSeqOut = self.gazeSeqCNNDropout(gazeSeqOut)

        headSeq = headSeq.reshape(-1, self.headSeqLength, self.headFeatureNum)
        headSeq = headSeq.permute(0,2,1)
        headSeqOut = self.headSeqCNN1D(headSeq)
        headSeqOut = headSeqOut.reshape(-1, self.headSeqCNN1D_outputSize)
        headSeqOut = self.headSeqCNNDropout(headSeqOut)

        objectSeq = objectSeq.reshape(-1, self.objectSeqLength, self.objectFeatureNum)
        objectSeq = objectSeq.permute(0,2,1)
        objectSeqOut = self.objectSeqCNN1D(objectSeq)
        objectSeqOut = objectSeqOut.reshape(-1, self.objectSeqCNN1D_outputSize)
        objectSeqOut = self.objectSeqCNNDropout(objectSeqOut)

        seqOut = torch.cat((headSeqOut, objectSeqOut), 1)

        saliencyMap = saliencyMap.reshape(saliencyMap.size(0), self.saliencyNum, self.saliencyWidth, self.saliencyWidth)
        saliencyFeatures = self.SalCNN(saliencyMap)
        saliencyOut = saliencyFeatures.reshape(saliencyMap.size(0), -1)
        saliencyOut = self.SalCNNDropout(saliencyOut)

        prdInput = torch.cat((seqOut, saliencyOut), 1)
        prdOut = self.PrdFC(prdInput)
        out = prdOut.mm(self.cluster)
        out = out + currentGaze
        return out

    def forward(self, x):
        out = self.forward1(x)
        return out

class FixationNet_without_Head(nn.Module):
    def __init__(self, gazeSeqSize, headSeqSize, objectSeqSize, saliencySize, clusterPath):
        super().__init__()
        # the input params
        self.gazeSeqSize = gazeSeqSize
        self.headSeqSize = headSeqSize
        self.objectSeqSize = objectSeqSize
        self.saliencySize = saliencySize

        print('gazeSeqSize: {}'.format(self.gazeSeqSize))
        print('headSeqSize: {}'.format(self.headSeqSize))
        print('taskObjectSeqSize: {}'.format(self.objectSeqSize))
        print('saliencySize: {}'.format(self.saliencySize))

        # preset params
        self.gazeFeatureNum = 2
        self.gazeSeqLength = int(self.gazeSeqSize/self.gazeFeatureNum)
        self.headFeatureNum = 2
        self.headSeqLength = int(self.headSeqSize/self.headFeatureNum)
        self.objectFeatureNum = 12
        self.objectSeqLength = int(self.objectSeqSize/self.objectFeatureNum)
        self.saliencyWidth = 24
        self.saliencyNum = int(self.saliencySize/(self.saliencyWidth*self.saliencyWidth))

        # load the cluster centers
        self.cluster = torch.from_numpy(np.load(clusterPath)).float().to(device)
        self.clusterSize = self.cluster.shape[0]
        print('Cluster Size: {}'.format(self.clusterSize))

        # GazeSeqCNN1D Module
        gazeSeqCNN1D_outChannels1 = 32
        gazeSeqCNN1D_poolingRate1 = 2
        gazeSeqCNN1D_kernelSize1 = 1
        gazeSeqCNN1D_featureSize1 = floor((self.gazeSeqLength - gazeSeqCNN1D_kernelSize1 + 1)/gazeSeqCNN1D_poolingRate1)
        self.gazeSeqCNN1D_outputSize = gazeSeqCNN1D_featureSize1 * gazeSeqCNN1D_outChannels1
        self.gazeSeqCNN1D = nn.Sequential(
            nn.Conv1d(in_channels=self.gazeFeatureNum, out_channels=gazeSeqCNN1D_outChannels1,kernel_size=gazeSeqCNN1D_kernelSize1),
            nn.BatchNorm1d(gazeSeqCNN1D_outChannels1),
            nn.ReLU(),
            #nn.LeakyReLU(),
            #nn.Sigmoid(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.MaxPool1d(gazeSeqCNN1D_poolingRate1),
             )

        gazeSeqCNN_dropoutRate = 0.0
        self.gazeSeqCNNDropout = nn.Sequential(
            nn.Dropout(p = gazeSeqCNN_dropoutRate)
             )

        # HeadSeqCNN1D Module
        headSeqCNN1D_outChannels1 = 64
        headSeqCNN1D_poolingRate1 = 2
        headSeqCNN1D_kernelSize1 = 1
        headSeqCNN1D_featureSize1 = floor((self.headSeqLength - headSeqCNN1D_kernelSize1 + 1)/headSeqCNN1D_poolingRate1)
        self.headSeqCNN1D_outputSize = headSeqCNN1D_featureSize1 * headSeqCNN1D_outChannels1
        self.headSeqCNN1D = nn.Sequential(
            nn.Conv1d(in_channels=self.headFeatureNum, out_channels=headSeqCNN1D_outChannels1,kernel_size=headSeqCNN1D_kernelSize1),
            nn.BatchNorm1d(headSeqCNN1D_outChannels1),
            nn.ReLU(),
            #nn.LeakyReLU(),
            #nn.Sigmoid(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.MaxPool1d(headSeqCNN1D_poolingRate1),
            #nn.Conv1d(in_channels=headSeqCNN1D_outChannels1, out_channels=headSeqCNN1D_outChannels2,kernel_size=headSeqCNN1D_kernelSize2),
            #nn.BatchNorm1d(headSeqCNN1D_outChannels2),
            #nn.Sigmoid(),
            #nn.MaxPool1d(headSeqCNN1D_poolingRate2),
             )

        headSeqCNN_dropoutRate = 0.0
        self.headSeqCNNDropout = nn.Sequential(
            nn.Dropout(p = headSeqCNN_dropoutRate)
             )

        # objectSeqCNN1D Module
        objectSeqCNN1D_outChannels1 = 64
        objectSeqCNN1D_poolingRate1 = 2
        objectSeqCNN1D_kernelSize1 = 1
        objectSeqCNN1D_featureSize1 = floor((self.objectSeqLength - objectSeqCNN1D_kernelSize1 + 1)/objectSeqCNN1D_poolingRate1)
        objectSeqCNN1D_outChannels2 = 32
        objectSeqCNN1D_poolingRate2 = 2
        objectSeqCNN1D_kernelSize2 = 1
        objectSeqCNN1D_featureSize2 = floor((objectSeqCNN1D_featureSize1 - objectSeqCNN1D_kernelSize2 + 1)/objectSeqCNN1D_poolingRate2)
        self.objectSeqCNN1D_outputSize = objectSeqCNN1D_featureSize2 * objectSeqCNN1D_outChannels2
        self.objectSeqCNN1D = nn.Sequential(
            nn.Conv1d(in_channels=self.objectFeatureNum, out_channels=objectSeqCNN1D_outChannels1,kernel_size=objectSeqCNN1D_kernelSize1),
            nn.BatchNorm1d(objectSeqCNN1D_outChannels1),
            nn.ReLU(),
            #nn.Sigmoid(),
            #nn.LeakyReLU(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.MaxPool1d(objectSeqCNN1D_poolingRate1),
            nn.Conv1d(in_channels=objectSeqCNN1D_outChannels1, out_channels=objectSeqCNN1D_outChannels2,kernel_size=objectSeqCNN1D_kernelSize2),
            nn.BatchNorm1d(objectSeqCNN1D_outChannels2),
            nn.ReLU(),
            #nn.Sigmoid(),
            nn.MaxPool1d(objectSeqCNN1D_poolingRate2),
             )

        objectSeqCNN_dropoutRate = 0.0
        self.objectSeqCNNDropout = nn.Sequential(
            nn.Dropout(p = objectSeqCNN_dropoutRate)
             )

        # SalCNN Module
        salCNN_outChannels1 = 8
        salCNN_poolingRate1 = 2
        salCNN_kernelSize1 = 1
        salCNN_padding1 = int((salCNN_kernelSize1-1)/2)
        salCNN_imageSize1 = floor((self.saliencyWidth - salCNN_kernelSize1 + 2*salCNN_padding1 + 1)/salCNN_poolingRate1)
        self.salCNN_outputSize = salCNN_imageSize1 * salCNN_imageSize1 * salCNN_outChannels1
        self.SalCNN = nn.Sequential(
            nn.Conv2d(self.saliencyNum, salCNN_outChannels1, kernel_size=salCNN_kernelSize1, padding=salCNN_padding1),
            nn.BatchNorm2d(salCNN_outChannels1),
            nn.ReLU(),
            #nn.Sigmoid(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.MaxPool2d(kernel_size=salCNN_poolingRate1),
             )

        salCNN_dropoutRate = 0.5
        self.SalCNNDropout = nn.Sequential(
            nn.Dropout(p = salCNN_dropoutRate)
             )

        # PrdFC Module
        prdFC_inputSize = self.gazeSeqCNN1D_outputSize + self.objectSeqCNN1D_outputSize + self.salCNN_outputSize
        prdFC_linearSize1 = 128
        prdFC_linearSize2 = self.clusterSize
        prdFC_dropoutRate = 0.5
        # the prediction fc layer for DGaze.
        self.PrdFC = nn.Sequential(
            nn.Linear(prdFC_inputSize, prdFC_linearSize1),
            nn.BatchNorm1d(prdFC_linearSize1),
            nn.ReLU(),
            #nn.Sigmoid(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.Dropout(p = prdFC_dropoutRate),
            nn.Linear(prdFC_linearSize1, prdFC_linearSize2),
            nn.Softmax(dim = 1),
             )

    def forward1(self, x):
        currentGaze = x[:, 0:2]
        index = self.gazeSeqSize
        gazeSeq = x[:, 0:index]
        #headSeq = x[:, index: index+self.headSeqSize]
        index += self.headSeqSize
        objectSeq = x[:, index: index+self.objectSeqSize]
        index += self.objectSeqSize
        saliencyMap = x[:, index: index+self.saliencySize]

        gazeSeq = gazeSeq.reshape(-1, self.gazeSeqLength, self.gazeFeatureNum)
        gazeSeq = gazeSeq.permute(0,2,1)
        gazeSeqOut = self.gazeSeqCNN1D(gazeSeq)
        gazeSeqOut = gazeSeqOut.reshape(-1, self.gazeSeqCNN1D_outputSize)
        gazeSeqOut = self.gazeSeqCNNDropout(gazeSeqOut)

        #headSeq = headSeq.reshape(-1, self.headSeqLength, self.headFeatureNum)
        #headSeq = headSeq.permute(0,2,1)
        #headSeqOut = self.headSeqCNN1D(headSeq)
        #headSeqOut = headSeqOut.reshape(-1, self.headSeqCNN1D_outputSize)
        #headSeqOut = self.headSeqCNNDropout(headSeqOut)

        objectSeq = objectSeq.reshape(-1, self.objectSeqLength, self.objectFeatureNum)
        objectSeq = objectSeq.permute(0,2,1)
        objectSeqOut = self.objectSeqCNN1D(objectSeq)
        objectSeqOut = objectSeqOut.reshape(-1, self.objectSeqCNN1D_outputSize)
        objectSeqOut = self.objectSeqCNNDropout(objectSeqOut)

        seqOut = torch.cat((gazeSeqOut, objectSeqOut), 1)

        saliencyMap = saliencyMap.reshape(saliencyMap.size(0), self.saliencyNum, self.saliencyWidth, self.saliencyWidth)
        saliencyFeatures = self.SalCNN(saliencyMap)
        saliencyOut = saliencyFeatures.reshape(saliencyMap.size(0), -1)
        saliencyOut = self.SalCNNDropout(saliencyOut)

        prdInput = torch.cat((seqOut, saliencyOut), 1)
        prdOut = self.PrdFC(prdInput)
        out = prdOut.mm(self.cluster)
        out = out + currentGaze
        return out

    def forward(self, x):
        out = self.forward1(x)
        return out

class FixationNet_without_Cluster(nn.Module):
    def __init__(self, gazeSeqSize, headSeqSize, objectSeqSize, saliencySize, clusterPath):
        super().__init__()
        # the input params
        self.gazeSeqSize = gazeSeqSize
        self.headSeqSize = headSeqSize
        self.objectSeqSize = objectSeqSize
        self.saliencySize = saliencySize

        print('gazeSeqSize: {}'.format(self.gazeSeqSize))
        print('headSeqSize: {}'.format(self.headSeqSize))
        print('taskObjectSeqSize: {}'.format(self.objectSeqSize))
        print('saliencySize: {}'.format(self.saliencySize))

        # preset params
        self.gazeFeatureNum = 2
        self.gazeSeqLength = int(self.gazeSeqSize/self.gazeFeatureNum)
        self.headFeatureNum = 2
        self.headSeqLength = int(self.headSeqSize/self.headFeatureNum)
        self.objectFeatureNum = 12
        self.objectSeqLength = int(self.objectSeqSize/self.objectFeatureNum)
        self.saliencyWidth = 24
        self.saliencyNum = int(self.saliencySize/(self.saliencyWidth*self.saliencyWidth))

        # GazeSeqCNN1D Module
        gazeSeqCNN1D_outChannels1 = 32
        gazeSeqCNN1D_poolingRate1 = 2
        gazeSeqCNN1D_kernelSize1 = 1
        gazeSeqCNN1D_featureSize1 = floor((self.gazeSeqLength - gazeSeqCNN1D_kernelSize1 + 1)/gazeSeqCNN1D_poolingRate1)
        self.gazeSeqCNN1D_outputSize = gazeSeqCNN1D_featureSize1 * gazeSeqCNN1D_outChannels1
        self.gazeSeqCNN1D = nn.Sequential(
            nn.Conv1d(in_channels=self.gazeFeatureNum, out_channels=gazeSeqCNN1D_outChannels1,kernel_size=gazeSeqCNN1D_kernelSize1),
            nn.BatchNorm1d(gazeSeqCNN1D_outChannels1),
            nn.ReLU(),
            #nn.LeakyReLU(),
            #nn.Sigmoid(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.MaxPool1d(gazeSeqCNN1D_poolingRate1),
             )

        gazeSeqCNN_dropoutRate = 0.0
        self.gazeSeqCNNDropout = nn.Sequential(
            nn.Dropout(p = gazeSeqCNN_dropoutRate)
             )

        # HeadSeqCNN1D Module
        headSeqCNN1D_outChannels1 = 64
        headSeqCNN1D_poolingRate1 = 2
        headSeqCNN1D_kernelSize1 = 1
        headSeqCNN1D_featureSize1 = floor((self.headSeqLength - headSeqCNN1D_kernelSize1 + 1)/headSeqCNN1D_poolingRate1)
        self.headSeqCNN1D_outputSize = headSeqCNN1D_featureSize1 * headSeqCNN1D_outChannels1
        self.headSeqCNN1D = nn.Sequential(
            nn.Conv1d(in_channels=self.headFeatureNum, out_channels=headSeqCNN1D_outChannels1,kernel_size=headSeqCNN1D_kernelSize1),
            nn.BatchNorm1d(headSeqCNN1D_outChannels1),
            nn.ReLU(),
            #nn.LeakyReLU(),
            #nn.Sigmoid(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.MaxPool1d(headSeqCNN1D_poolingRate1),
            #nn.Conv1d(in_channels=headSeqCNN1D_outChannels1, out_channels=headSeqCNN1D_outChannels2,kernel_size=headSeqCNN1D_kernelSize2),
            #nn.BatchNorm1d(headSeqCNN1D_outChannels2),
            #nn.Sigmoid(),
            #nn.MaxPool1d(headSeqCNN1D_poolingRate2),
             )

        headSeqCNN_dropoutRate = 0.0
        self.headSeqCNNDropout = nn.Sequential(
            nn.Dropout(p = headSeqCNN_dropoutRate)
             )

        # objectSeqCNN1D Module
        objectSeqCNN1D_outChannels1 = 64
        objectSeqCNN1D_poolingRate1 = 2
        objectSeqCNN1D_kernelSize1 = 1
        objectSeqCNN1D_featureSize1 = floor((self.objectSeqLength - objectSeqCNN1D_kernelSize1 + 1)/objectSeqCNN1D_poolingRate1)
        objectSeqCNN1D_outChannels2 = 32
        objectSeqCNN1D_poolingRate2 = 2
        objectSeqCNN1D_kernelSize2 = 1
        objectSeqCNN1D_featureSize2 = floor((objectSeqCNN1D_featureSize1 - objectSeqCNN1D_kernelSize2 + 1)/objectSeqCNN1D_poolingRate2)
        self.objectSeqCNN1D_outputSize = objectSeqCNN1D_featureSize2 * objectSeqCNN1D_outChannels2
        self.objectSeqCNN1D = nn.Sequential(
            nn.Conv1d(in_channels=self.objectFeatureNum, out_channels=objectSeqCNN1D_outChannels1,kernel_size=objectSeqCNN1D_kernelSize1),
            nn.BatchNorm1d(objectSeqCNN1D_outChannels1),
            nn.ReLU(),
            #nn.Sigmoid(),
            #nn.LeakyReLU(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.MaxPool1d(objectSeqCNN1D_poolingRate1),
            nn.Conv1d(in_channels=objectSeqCNN1D_outChannels1, out_channels=objectSeqCNN1D_outChannels2,kernel_size=objectSeqCNN1D_kernelSize2),
            nn.BatchNorm1d(objectSeqCNN1D_outChannels2),
            nn.ReLU(),
            #nn.Sigmoid(),
            nn.MaxPool1d(objectSeqCNN1D_poolingRate2),
             )

        objectSeqCNN_dropoutRate = 0.0
        self.objectSeqCNNDropout = nn.Sequential(
            nn.Dropout(p = objectSeqCNN_dropoutRate)
             )

        # SalCNN Module
        salCNN_outChannels1 = 8
        salCNN_poolingRate1 = 2
        salCNN_kernelSize1 = 1
        salCNN_padding1 = int((salCNN_kernelSize1-1)/2)
        salCNN_imageSize1 = floor((self.saliencyWidth - salCNN_kernelSize1 + 2*salCNN_padding1 + 1)/salCNN_poolingRate1)
        self.salCNN_outputSize = salCNN_imageSize1 * salCNN_imageSize1 * salCNN_outChannels1
        self.SalCNN = nn.Sequential(
            nn.Conv2d(self.saliencyNum, salCNN_outChannels1, kernel_size=salCNN_kernelSize1, padding=salCNN_padding1),
            nn.BatchNorm2d(salCNN_outChannels1),
            nn.ReLU(),
            #nn.Sigmoid(),
            #nn.Tanh(),
            #nn.Softplus(),
            nn.MaxPool2d(kernel_size=salCNN_poolingRate1),
             )

        salCNN_dropoutRate = 0.5
        self.SalCNNDropout = nn.Sequential(
            nn.Dropout(p = salCNN_dropoutRate)
             )

        # PrdFC Module
        prdFC_inputSize = self.gazeSeqCNN1D_outputSize + self.headSeqCNN1D_outputSize + self.objectSeqCNN1D_outputSize + self.salCNN_outputSize
        prdFC_linearSize1 = 128
        prdFC_linearSize2 = 2
        #prdFC_dropoutRate = 0.0
        # the prediction fc layer for DGaze.
        self.PrdFC = nn.Sequential(
            nn.Linear(prdFC_inputSize, prdFC_linearSize1),
            nn.BatchNorm1d(prdFC_linearSize1),
            nn.ReLU(),
            #nn.Sigmoid(),
            #nn.Tanh(),
            #nn.Softplus(),
            #nn.Dropout(p = prdFC_dropoutRate),
            nn.Linear(prdFC_linearSize1, prdFC_linearSize2)
             )

    def forward1(self, x):
        index = self.gazeSeqSize
        gazeSeq = x[:, 0:index]
        headSeq = x[:, index: index+self.headSeqSize]
        index += self.headSeqSize
        objectSeq = x[:, index: index+self.objectSeqSize]
        index += self.objectSeqSize
        saliencyMap = x[:, index: index+self.saliencySize]

        gazeSeq = gazeSeq.reshape(-1, self.gazeSeqLength, self.gazeFeatureNum)
        gazeSeq = gazeSeq.permute(0,2,1)
        gazeSeqOut = self.gazeSeqCNN1D(gazeSeq)
        gazeSeqOut = gazeSeqOut.reshape(-1, self.gazeSeqCNN1D_outputSize)
        gazeSeqOut = self.gazeSeqCNNDropout(gazeSeqOut)

        headSeq = headSeq.reshape(-1, self.headSeqLength, self.headFeatureNum)
        headSeq = headSeq.permute(0,2,1)
        headSeqOut = self.headSeqCNN1D(headSeq)
        headSeqOut = headSeqOut.reshape(-1, self.headSeqCNN1D_outputSize)
        headSeqOut = self.headSeqCNNDropout(headSeqOut)

        objectSeq = objectSeq.reshape(-1, self.objectSeqLength, self.objectFeatureNum)
        objectSeq = objectSeq.permute(0,2,1)
        objectSeqOut = self.objectSeqCNN1D(objectSeq)
        objectSeqOut = objectSeqOut.reshape(-1, self.objectSeqCNN1D_outputSize)
        objectSeqOut = self.objectSeqCNNDropout(objectSeqOut)

        seqOut = torch.cat((gazeSeqOut, headSeqOut), 1)
        seqOut = torch.cat((seqOut, objectSeqOut), 1)

        saliencyMap = saliencyMap.reshape(saliencyMap.size(0), self.saliencyNum, self.saliencyWidth, self.saliencyWidth)
        saliencyFeatures = self.SalCNN(saliencyMap)
        saliencyOut = saliencyFeatures.reshape(saliencyMap.size(0), -1)
        saliencyOut = self.SalCNNDropout(saliencyOut)

        prdInput = torch.cat((seqOut, saliencyOut), 1)
        prdOut = self.PrdFC(prdInput)
        out = prdOut
        return out

    def forward(self, x):
        out = self.forward1(x)
        return out

max_stars_repo_name: CraneHzm/FixationNet

max_issues_repo_name: CraneHzm/FixationNet

max_forks_repo_name: CraneHzm/FixationNet


text:
# -*- coding: utf-8 -*-
#  Copyright (c) 2020. Distributed under the terms of the MIT License.
from argparse import Namespace
from pathlib import Path

import numpy as np
from pydefect.analyzer.calc_results import CalcResults
from pydefect.analyzer.defect_structure_info import DefectStructureInfo
from pydefect.analyzer.unitcell import Unitcell
from pydefect.cli.vasp.main_vasp_functions import make_unitcell, \
    make_competing_phase_dirs, make_composition_energies, make_defect_entries, \
    make_calc_results, make_band_edge_orb_infos_and_eigval_plot, \
    make_perfect_band_edge_state, make_local_extrema
from pydefect.input_maker.defect import SimpleDefect
from pydefect.input_maker.defect_entry import DefectEntry
from pydefect.input_maker.defect_set import DefectSet
from pymatgen.core import Composition, Structure
from pymatgen.io.vasp import Vasprun, Outcar, VolumetricData
from vise.defaults import defaults

def test_make_unitcell(tmpdir, mocker):
    tmpdir.chdir()
    vasprun_band_mock = mocker.Mock(spec=Vasprun, autospec=True)
    outcar_band_mock = mocker.Mock(spec=Outcar, autospec=True)
    outcar_dielectric_mock = mocker.Mock(spec=Outcar, autospec=True)
    args = Namespace(vasprun_band=vasprun_band_mock,
                     outcar_band=outcar_band_mock,
                     outcar_dielectric_clamped=outcar_dielectric_mock,
                     outcar_dielectric_ionic=outcar_dielectric_mock,
                     name="H")

    mock = mocker.patch(
        "pydefect.cli.vasp.main_vasp_functions.make_unitcell_from_vasp")
    mock.return_value = Unitcell(system="H",
                                 vbm=1.0,
                                 cbm=2.0,
                                 ele_dielectric_const=np.eye(3).tolist(),
                                 ion_dielectric_const=np.eye(3).tolist())
    make_unitcell(args)
    mock.assert_called_once_with(
        vasprun_band=vasprun_band_mock,
        outcar_band=outcar_band_mock,
        outcar_dielectric_clamped=outcar_dielectric_mock,
        outcar_dielectric_ionic=outcar_dielectric_mock,
        system_name="H")

def test_make_competing_phase_dirs(mocker):
    args = Namespace(elements=["Mg", "O"],
                     e_above_hull=0.1)
    mock = mocker.patch("pydefect.cli.vasp.main_vasp_functions.MpQuery")
    mock_make = mocker.patch(
        "pydefect.cli.vasp.main_vasp_functions.make_poscars_from_query")
    make_competing_phase_dirs(args)
    mock.assert_called_once_with(element_list=args.elements,
                                 e_above_hull=args.e_above_hull)
    mock_make.assert_called_once_with(
        materials_query=mock.return_value.materials, path=Path.cwd())

def test_make_composition_energies(mocker, tmpdir):
    def side_effect(key):
        mock_outcar = mocker.Mock()
        if key == Path("Mg") / defaults.outcar:
            mock_outcar.final_energy = -10
        elif key == Path("O") / defaults.outcar:
            mock_outcar.final_energy = -20
        else:
            raise ValueError
        return mock_outcar

    def side_effect_structure(key):
        mock_structure = mocker.Mock()
        if key == Path("Mg") / defaults.contcar:
            mock_structure.composition = Composition("Mg2")
        elif key == Path("O") / defaults.contcar:
            mock_structure.composition = Composition("O2")
        else:
            raise ValueError
        return mock_structure

    tmpdir.chdir()
    print(tmpdir)
    mock = mocker.patch("pydefect.cli.vasp.main_vasp_functions.Outcar",
                        side_effect=side_effect)
    mock_str = mocker.patch("pydefect.cli.vasp.main_vasp_functions.Structure.from_file",
                            side_effect=side_effect_structure)
    args = Namespace(yaml_file=None, dirs=[Path("Mg"), Path("O")])
    make_composition_energies(args)
    actual = Path("composition_energies.yaml").read_text()
    expected = """Mg2:
  energy: -10.0
  source: local
O2:
  energy: -20.0
  source: local
"""
    assert actual == expected

def test_make_local_extrema(tmpdir, mocker, simple_cubic):
    print(tmpdir)
    tmpdir.chdir()

    mock_params = mocker.patch("pydefect.cli.vasp.main_vasp_functions.VolumetricDataAnalyzeParams")
    mock_make_extrema = mocker.patch("pydefect.cli.vasp.main_vasp_functions.make_local_extrema_from_volumetric_data")
    volumetric_data = VolumetricData(simple_cubic,
                                     data={"total": np.array([[[0.0]]])})
    args = Namespace(volumetric_data=[volumetric_data],
                     find_max=True,
                     info="a",
                     threshold_frac=None,
                     threshold_abs=None,
                     min_dist=0.1,
                     tol=0.2,
                     radius=0.3)
    make_local_extrema(args)
    mock_params.assert_called_once_with(None, None, 0.1, 0.2, 0.3)
    mock_make_extrema.assert_called_once_with(volumetric_data=volumetric_data,
                                              params=mock_params.return_value,
                                              info="a",
                                              find_min=False)
    mock_make_extrema.return_value.to_json_file.assert_called_once_with()

def test_make_defect_entries(tmpdir, supercell_info):
    tmpdir.chdir()
    supercell_info.to_json_file()
    defect_set = DefectSet({SimpleDefect(None, "He1", [-1, 0])})
    defect_set.to_yaml()
    args = Namespace()
    make_defect_entries(args)
    names = {str(name) for name in Path(".").glob("*")}
    assert names == {'Va_He1_-1', 'defect_in.yaml', 'perfect', 'Va_He1_0',
                     'supercell_info.json'}

    perfect_structure = Structure.from_file(str(Path("perfect") / "POSCAR"))
    assert perfect_structure == supercell_info.structure

    file_names = {str(file_name.name)
                  for file_name in Path("Va_He1_-1").glob("*")}
    assert file_names == {"POSCAR", "defect_entry.json", "prior_info.yaml"}

    expected = """charge: -1
"""
    assert Path("Va_He1_-1/prior_info.yaml").read_text() == expected

def test_make_calc_results(tmpdir, mocker):
    tmpdir.chdir()
    mock = mocker.patch(
        "pydefect.cli.vasp.main_vasp_functions.make_calc_results_from_vasp")
    mock_vasprun = mocker.patch("pydefect.cli.vasp.main_vasp_functions.Vasprun")
    mock_outcar = mocker.patch("pydefect.cli.vasp.main_vasp_functions.Outcar")
    mock_calc_results = mocker.Mock(spec=CalcResults)
    mock.return_value = mock_calc_results
    args = Namespace(dirs=[Path("a")])
    make_calc_results(args)

    mock_vasprun.assert_called_with(Path("a") / defaults.vasprun)
    mock_outcar.assert_called_with(Path("a") / defaults.outcar)
    mock.assert_called_with(vasprun=mock_vasprun.return_value,
                            outcar=mock_outcar.return_value)
    mock_calc_results.to_json_file.assert_called_with("a/calc_results.json")

def test_make_perfect_band_edge_state(mocker):
    mock_vasprun = mocker.patch("pydefect.cli.vasp.main_vasp_functions.Vasprun")
    mock_procar = mocker.patch("pydefect.cli.vasp.main_vasp_functions.Procar")
    mock_outcar = mocker.patch("pydefect.cli.vasp.main_vasp_functions.Outcar")
    mock_make_perf_be_state = mocker.patch(
        "pydefect.cli.vasp.main_vasp_functions.make_perfect_band_edge_state_from_vasp")

    args = Namespace(dir=Path("perfect"))
    make_perfect_band_edge_state(args)

    mock_vasprun.assert_called_with(Path("perfect") / defaults.vasprun)
    mock_procar.assert_called_with(Path("perfect") / defaults.procar)
    mock_outcar.assert_called_with(Path("perfect") / defaults.outcar)
    mock_make_perf_be_state.assert_called_with(mock_procar.return_value,
                                               mock_vasprun.return_value,
                                               mock_outcar.return_value)

def test_make_band_edge_orb_infos_and_eigval_plot(mocker):
    mock_procar = mocker.patch("pydefect.cli.vasp.main_vasp_functions.Procar")
    mock_vasprun = mocker.patch("pydefect.cli.vasp.main_vasp_functions.Vasprun")

    mock_p_state = mocker.Mock()
    mock_p_state.vbm_info.energy = 10
    mock_p_state.cbm_info.energy = 20

    mock_defect_entry = mocker.Mock(spec=DefectEntry, autospec=True)

    mock_structure_info = mocker.Mock(spec=DefectStructureInfo, autospec=True)

    mock_make_orbital_infos = mocker.patch(
        "pydefect.cli.vasp.main_vasp_functions.make_band_edge_orbital_infos")
    mock_eigval_plotter = mocker.patch(
        "pydefect.cli.vasp.main_vasp_functions.EigenvalueMplPlotter")

    def side_effect(key):
        if str(key) == "Va_O1_2/defect_entry.json":
            mock_defect_entry.name = "Va_O1"
            mock_defect_entry.charge = 2
            return mock_defect_entry
        elif str(key) == "Va_O1_2/defect_structure_info.json":
            return mock_structure_info
        else:
            raise ValueError

    mock_loadfn = mocker.patch("pydefect.cli.vasp.main_vasp_functions.loadfn",
                               side_effect=side_effect)

    args = Namespace(dirs=[Path("Va_O1_2")],
                     p_state=mock_p_state,
                     y_range=[0.0, 1.0])
    make_band_edge_orb_infos_and_eigval_plot(args)

    mock_procar.assert_called_with(Path("Va_O1_2") / defaults.procar)
    mock_vasprun.assert_called_with(Path("Va_O1_2") / defaults.vasprun)

    mock_make_orbital_infos.assert_called_with(
        mock_procar.return_value, mock_vasprun.return_value, 10, 20,
        mock_structure_info)

    mock_loadfn.assert_any_call(Path("Va_O1_2") / "defect_entry.json")

    mock_eigval_plotter.assert_called_with(
        title="Va_O1",
        band_edge_orb_infos=mock_make_orbital_infos.return_value,
        supercell_vbm=10,
        supercell_cbm=20,
        y_range=[0.0, 1.0])

ext: py

max_stars_repo_name: kumagai-group/pydefect

max_issues_repo_name: kumagai-group/pydefect

max_forks_repo_name: kumagai-group/pydefect

reason: import numpy


text:
import os
import numpy as np
from .read_buffer_offline import read_buffer_offline_data, read_buffer_offline_events, read_buffer_offline_header
from mindaffectBCI.decoder.utils import block_randomize, butter_sosfilt, upsample_codebook, lab2ind, window_axis

trigger_event='stimulus.note.play' # the actual times the user hit the button

def load_brainsonfire(datadir, sessdir=None, sessfn=None, fs_out=60, stopband=((45,65),(0,1),(25,-1)), subtriallen=10, nvirt=20, chIdx=slice(64), verb=2):

    # load the data file
    Xfn = datadir
    if sessdir:
        Xfn = os.path.join(Xfn, sessdir)
    if sessfn:
        Xfn = os.path.join(Xfn, sessfn)
    sessdir = os.path.dirname(Xfn)

    if verb > 1: print("Loading header")
    hdr=read_buffer_offline_header(Xfn)
    if verb > 1: print("Loading data")
    X = read_buffer_offline_data(Xfn,hdr) # (nsamp,nch)
    if verb > 1: print("Loading events")
    evts=read_buffer_offline_events(Xfn)

    fs = hdr.fs
    ch_names = hdr.labels

    if chIdx is not None:
        X = X [...,chIdx]
        ch_names = ch_names[chIdx] if ch_names is not None else None

    # pre-resample to save memory
    rsrate = int(fs//120)
    if rsrate > 1:
        if verb > 0: print("Pre-re-sample by {}: {}->{}Hz".format(rsrate,fs,fs/rsrate))
        X = X [::rsrate,:]
        for e in evts:
            e.sample = e.sample/rsrate
        fs = fs/rsrate

    if verb > 0: print("X={} @{}Hz".format(X.shape,fs),flush=True)

    # extract the trigger info
    trigevts = [e for e in evts if e.type.lower() == trigger_event]
    trig_samp= np.array([e.sample for e in trigevts],dtype=int)
    trig_val = [e.value for e in trigevts]
    trig_ind, lab2class = lab2ind(trig_val) # convert to indicator (ntrig,ncls)
    # up-sample to stim rate
    Y = np.zeros((X.shape[0],trig_ind.shape[-1]),dtype=bool)
    Y[trig_samp,:] = trig_ind
    if verb > 0:
            print("Y={}".format(Y.shape))

    # BODGE: trim to useful data range
    if .1 < (trig_samp[0]-fs)/X.shape[0] or (trig_samp[-1]+fs)/X.shape[0] < .9:
        if verb>0 : print('Trimming range: {}-{}s'.format(trig_samp[0]/fs,trig_samp[-1]/fs))
        # limit to the useful data range
        rng = slice(int(trig_samp[0]-fs), int(trig_samp[-1]+fs))
        X = X[rng, :]
        Y = Y[rng, ...]
        if verb > 0: print("X={}".format(X.shape))
        if verb > 0: print("Y={}".format(Y.shape))

    # preprocess -> spectral filter, in continuous time!
    if stopband is not None:
        if verb > 0:
            print("preFilter: {}Hz".format(stopband))
        X, _, _ = butter_sosfilt(X,stopband,fs)

    # preprocess -> downsample
    resamprate = int(fs/fs_out)
    if resamprate > 1:
        if verb > 0:
            print("resample by {}: {}->{}Hz".format(resamprate, fs, fs/resamprate))
        X = X[..., ::resamprate, :] # decimate X (trl, samp, d)
        # re-sample Y, being sure to keep any events in the re-sample window
        Y = window_axis(Y,winsz=resamprate,step=resamprate,axis=-2) # (trl, samp, win, e)
        Y = np.max(Y,axis=-2) # (trl,samp,e)  N.B. use max so don't loose single sample events
        fs = fs/resamprate

    # make virtual targets
    Y = Y[:,np.newaxis,:] # (nsamp,1,e)
    Y_virt = block_randomize(Y, nvirt, axis=-3) # (nsamp,nvirt,e)
    Y = np.concatenate((Y, Y_virt), axis=-2) # (nsamp,1+nvirt,e)
    if verb > 0: print("Y={}".format(Y.shape))

    # cut into sub-trials
    nsubtrials = X.shape[0]/fs/subtriallen
    if nsubtrials > 1:
        winsz = int(X.shape[0]//nsubtrials)
        if verb > 0: print('subtrial winsz={}'.format(winsz))
        # slice into sub-trials
        X = window_axis(X,axis=0,winsz=winsz,step=winsz) # (trl,win,d)
        Y = window_axis(Y,axis=0,winsz=winsz,step=winsz) # (trl,win,nY)
        if verb > 0: print("X={}".format(X.shape))
        if verb > 0: print("Y={}".format(Y.shape))

    # make coords array for the meta-info about the dimensions of X
    coords = [None]*X.ndim
    coords[0] = {'name':'trial'}
    coords[1] = {'name':'time','unit':'ms', \
                 'coords':np.arange(X.shape[1])/fs, \
                 'fs':fs}
    coords[2] = {'name':'channel','coords':ch_names}
    # return data + metadata
    return (X, Y, coords)

def testcase():
    import sys

    if os.path.isdir('D:\external_data'):
        datadir = 'D:/own_experiments/'
    else:
        datadir = '/home/jadref/data/bci/own_experiments'
    sessfn = os.path.join(datadir,'motor_imagery/brainsonfire/brains_on_fire_online/subject01/raw_buffer/0001')
    # command-line, for testing
    if len(sys.argv) > 1:
        sessfn = sys.argv[1]

    from load_brainsonfire import load_brainsonfire
    print("Loading: {}".format(sessfn))
    oX, oY, coords = load_brainsonfire(sessfn, fs_out=60)
    times = coords[1]['coords']
    fs = coords[1]['fs']
    ch_names = coords[2]['coords']
    X=oX.copy()
    Y=oY.copy()

    print("X({}){}".format([c['name'] for c in coords],X.shape))
    print("Y={}".format(Y.shape))
    print("fs={}".format(fs))

    tau=fs*.3
    evtlabs = None
    times=np.arange(int(tau))/fs
    rank=10

    # visualize the dataset
    from stim2event import stim2event
    from updateSummaryStatistics import updateSummaryStatistics, plot_erp, plot_summary_statistics, idOutliers
    import matplotlib.pyplot as plt

    Cxx, Cxy, Cyy = updateSummaryStatistics(X, Y[...,0:1,:], tau=tau)

    plt.figure(1);
    print("summary stats")
    plot_summary_statistics(Cxx, Cxy, Cyy, evtlabs, times, ch_names)

    plt.figure(2);
    print("ERP")
    plot_erp(Cxy, ch_names=ch_names, evtlabs=evtlabs, times=times, plottype='plot', axis=-1)

    from model_fitting import MultiCCA
    from decodingCurveSupervised import decodingCurveSupervised
    cca = MultiCCA(tau=tau, evtlabs=evtlabs, rank=rank)
    scores = cca.cv_fit(X, Y)
    Fy = scores['estimator']
    print("Fy={}".format(Fy.shape))
    (_)=decodingCurveSupervised(Fy)

    # plot the solution
    from scoreStimulus import factored2full
    print("Plot Model")
    plt.figure(3)
    plot_erp(factored2full(cca.W_, cca.R_), ch_names=ch_names, evtlabs=evtlabs, times=times)
    #   plot Fy
    plt.figure(4)
    for ti in range(min(Fy.shape[0],25)):
        plt.subplot(5,5,ti+1)
        plt.imshow(np.cumsum(Fy[ti,:,:],axis=-2),aspect='auto')
    plt.show()

if __name__=="__main__":
    testcase()

max_issues_repo_name: CkiChen/pymindaffectBCI

max_forks_repo_name: CkiChen/pymindaffectBCI


text:
from numba import jit, jitclass, float64
import numpy as np

@jit(nopython=True)
def _B1_LT_085(x):
    # B1.LT_085: lt(0.85)
    if 0.0 == 0.0:
        return 1.0 if x < 0.85 else 0.0
    x1 = 0.85 - 0.0
    x2 = 0.85 + 0.0
    if x <= x1:
        return 1.0
    if x <= x2:
        return 1.0 - (x - x1) / (x2 - x1)
    return 0.0

@jit(nopython=True)
def _B1_GT_1(x):
    # B1.GT_1: gt(1.0)
    if 0.0 == 0.0:
        return 1.0 if x > 1.0 else 0.0
    x1 = 1.0 - 0.0
    x2 = 1.0 + 0.0
    if x <= x1:
        return 0.0
    if x <= x2:
        return (x - x1) / (x2 - x1)
    return 1.0

@jit(nopython=True)
def _B2_GT_0(x):
    # B2.GT_0: gt(0.0)
    if 0.0 == 0.0:
        return 1.0 if x > 0.0 else 0.0
    x1 = 0.0 - 0.0
    x2 = 0.0 + 0.0
    if x <= x1:
        return 0.0
    if x <= x2:
        return (x - x1) / (x2 - x1)
    return 1.0

@jit(nopython=True)
def _B3_LT_005(x):
    # B3.LT_005: lt(0.05)
    if 0.0 == 0.0:
        return 1.0 if x < 0.05 else 0.0
    x1 = 0.05 - 0.0
    x2 = 0.05 + 0.0
    if x <= x1:
        return 1.0
    if x <= x2:
        return 1.0 - (x - x1) / (x2 - x1)
    return 0.0

@jit(nopython=True)
def _B3_LT_01(x):
    # B3.LT_01: lt(0.1)
    if 0.0 == 0.0:
        return 1.0 if x < 0.1 else 0.0
    x1 = 0.1 - 0.0
    x2 = 0.1 + 0.0
    if x <= x1:
        return 1.0
    if x <= x2:
        return 1.0 - (x - x1) / (x2 - x1)
    return 0.0

@jit(nopython=True)
def _B3_LT_015(x):
    # B3.LT_015: lt(0.15)
    if 0.0 == 0.0:
        return 1.0 if x < 0.15 else 0.0
    x1 = 0.15 - 0.0
    x2 = 0.15 + 0.0
    if x <= x1:
        return 1.0
    if x <= x2:
        return 1.0 - (x - x1) / (x2 - x1)
    return 0.0

@jit(nopython=True)
def _B3_LT_02(x):
    # B3.LT_02: lt(0.2)
    if 0.0 == 0.0:
        return 1.0 if x < 0.2 else 0.0
    x1 = 0.2 - 0.0
    x2 = 0.2 + 0.0
    if x <= x1:
        return 1.0
    if x <= x2:
        return 1.0 - (x - x1) / (x2 - x1)
    return 0.0

@jit(nopython=True)
def _B4_NODATA(x):
    # B4.NODATA: eq(0.0)
    if 0.0 == 0.0:
        return 1.0 if x == 0.0 else 0.0
    x1 = 0.0 - 0.0
    x2 = 0.0
    x3 = 0.0 + 0.0
    if x <= x1:
        return 0.0
    if x <= x2:
        return (x - x1) / (x2 - x1)
    if x <= x3:
        return 1.0 - (x - x2) / (x3 - x2)
    return 0.0

@jit(nopython=True)
def _B5_LT_01(x):
    # B5.LT_01: lt(0.1)
    if 0.0 == 0.0:
        return 1.0 if x < 0.1 else 0.0
    x1 = 0.1 - 0.0
    x2 = 0.1 + 0.0
    if x <= x1:
        return 1.0
    if x <= x2:
        return 1.0 - (x - x1) / (x2 - x1)
    return 0.0

@jit(nopython=True)
def _B7_LT_05(x):
    # B7.LT_05: lt(0.5)
    if 0.0 == 0.0:
        return 1.0 if x < 0.5 else 0.0
    x1 = 0.5 - 0.0
    x2 = 0.5 + 0.0
    if x <= x1:
        return 1.0
    if x <= x2:
        return 1.0 - (x - x1) / (x2 - x1)
    return 0.0

@jit(nopython=True)
def _B8_GT_0(x):
    # B8.GT_0: gt(0.0)
    if 0.0 == 0.0:
        return 1.0 if x > 0.0 else 0.0
    x1 = 0.0 - 0.0
    x2 = 0.0 + 0.0
    if x <= x1:
        return 0.0
    if x <= x2:
        return (x - x1) / (x2 - x1)
    return 1.0

@jit(nopython=True)
def _B8_LT_009(x):
    # B8.LT_009: lt(0.09)
    if 0.0 == 0.0:
        return 1.0 if x < 0.09 else 0.0
    x1 = 0.09 - 0.0
    x2 = 0.09 + 0.0
    if x <= x1:
        return 1.0
    if x <= x2:
        return 1.0 - (x - x1) / (x2 - x1)
    return 0.0

@jit(nopython=True)
def _B8_GT_033(x):
    # B8.GT_033: gt(0.33)
    if 0.0 == 0.0:
        return 1.0 if x > 0.33 else 0.0
    x1 = 0.33 - 0.0
    x2 = 0.33 + 0.0
    if x <= x1:
        return 0.0
    if x <= x2:
        return (x - x1) / (x2 - x1)
    return 1.0

@jit(nopython=True)
def _B8_GT_035(x):
    # B8.GT_035: gt(0.35)
    if 0.0 == 0.0:
        return 1.0 if x > 0.35 else 0.0
    x1 = 0.35 - 0.0
    x2 = 0.35 + 0.0
    if x <= x1:
        return 0.0
    if x <= x2:
        return (x - x1) / (x2 - x1)
    return 1.0

@jit(nopython=True)
def _B8_GT_04(x):
    # B8.GT_04: gt(0.4)
    if 0.0 == 0.0:
        return 1.0 if x > 0.4 else 0.0
    x1 = 0.4 - 0.0
    x2 = 0.4 + 0.0
    if x <= x1:
        return 0.0
    if x <= x2:
        return (x - x1) / (x2 - x1)
    return 1.0

@jit(nopython=True)
def _B8_GT_045(x):
    # B8.GT_045: gt(0.45)
    if 0.0 == 0.0:
        return 1.0 if x > 0.45 else 0.0
    x1 = 0.45 - 0.0
    x2 = 0.45 + 0.0
    if x <= x1:
        return 0.0
    if x <= x2:
        return (x - x1) / (x2 - x1)
    return 1.0

@jit(nopython=True)
def _B8_LT_085(x):
    # B8.LT_085: lt(0.85)
    if 0.0 == 0.0:
        return 1.0 if x < 0.85 else 0.0
    x1 = 0.85 - 0.0
    x2 = 0.85 + 0.0
    if x <= x1:
        return 1.0
    if x <= x2:
        return 1.0 - (x - x1) / (x2 - x1)
    return 0.0

@jit(nopython=True)
def _B16_GT_0(x):
    # B16.GT_0: gt(0.0)
    if 0.0 == 0.0:
        return 1.0 if x > 0.0 else 0.0
    x1 = 0.0 - 0.0
    x2 = 0.0 + 0.0
    if x <= x1:
        return 0.0
    if x <= x2:
        return (x - x1) / (x2 - x1)
    return 1.0

@jit(nopython=True)
def _B19_GT_015(x):
    # B19.GT_015: gt(0.15)
    if 0.0 == 0.0:
        return 1.0 if x > 0.15 else 0.0
    x1 = 0.15 - 0.0
    x2 = 0.15 + 0.0
    if x <= x1:
        return 0.0
    if x <= x2:
        return (x - x1) / (x2 - x1)
    return 1.0

@jit(nopython=True)
def _BSum_GT_011(x):
    # BSum.GT_011: gt(0.11)
    if 0.0 == 0.0:
        return 1.0 if x > 0.11 else 0.0
    x1 = 0.11 - 0.0
    x2 = 0.11 + 0.0
    if x <= x1:
        return 0.0
    if x <= x2:
        return (x - x1) / (x2 - x1)
    return 1.0

@jit(nopython=True)
def _BSum_GT_013(x):
    # BSum.GT_013: gt(0.13)
    if 0.0 == 0.0:
        return 1.0 if x > 0.13 else 0.0
    x1 = 0.13 - 0.0
    x2 = 0.13 + 0.0
    if x <= x1:
        return 0.0
    if x <= x2:
        return (x - x1) / (x2 - x1)
    return 1.0

@jit(nopython=True)
def _BSum_GT_016(x):
    # BSum.GT_016: gt(0.16)
    if 0.0 == 0.0:
        return 1.0 if x > 0.16 else 0.0
    x1 = 0.16 - 0.0
    x2 = 0.16 + 0.0
    if x <= x1:
        return 0.0
    if x <= x2:
        return (x - x1) / (x2 - x1)
    return 1.0

@jit(nopython=True)
def _Class_FALSE(x):
    # Class.FALSE: false()
    return 0.0

@jit(nopython=True)
def _Class_TRUE(x):
    # Class.TRUE: true()
    return 1.0

_InputsSpec = [
    ("b1", float64[:]),
    ("b2", float64[:]),
    ("b3", float64[:]),
    ("b4", float64[:]),
    ("b5", float64[:]),
    ("b6", float64[:]),
    ("b7", float64[:]),
    ("b8", float64[:]),
    ("b12", float64[:]),
    ("b13", float64[:]),
    ("b14", float64[:]),
    ("b15", float64[:]),
    ("b16", float64[:]),
    ("b19", float64[:]),
    ("b100", float64[:]),
    ("bsum", float64[:]),
]

@jitclass(_InputsSpec)
class Inputs:
    def __init__(self, size: int):
        self.b1 = np.zeros(size, dtype=np.float64)
        self.b2 = np.zeros(size, dtype=np.float64)
        self.b3 = np.zeros(size, dtype=np.float64)
        self.b4 = np.zeros(size, dtype=np.float64)
        self.b5 = np.zeros(size, dtype=np.float64)
        self.b6 = np.zeros(size, dtype=np.float64)
        self.b7 = np.zeros(size, dtype=np.float64)
        self.b8 = np.zeros(size, dtype=np.float64)
        self.b12 = np.zeros(size, dtype=np.float64)
        self.b13 = np.zeros(size, dtype=np.float64)
        self.b14 = np.zeros(size, dtype=np.float64)
        self.b15 = np.zeros(size, dtype=np.float64)
        self.b16 = np.zeros(size, dtype=np.float64)
        self.b19 = np.zeros(size, dtype=np.float64)
        self.b100 = np.zeros(size, dtype=np.float64)
        self.bsum = np.zeros(size, dtype=np.float64)

_OutputsSpec = [
    ("nodata", float64[:]),
    ("Wasser", float64[:]),
    ("Schill", float64[:]),
    ("Muschel", float64[:]),
    ("dense2", float64[:]),
    ("dense1", float64[:]),
    ("Strand", float64[:]),
    ("Sand", float64[:]),
    ("Misch", float64[:]),
    ("Misch2", float64[:]),
    ("Schlick", float64[:]),
    ("schlick_t", float64[:]),
    ("Wasser2", float64[:]),
]

@jitclass(_OutputsSpec)
class Outputs:
    def __init__(self, size: int):
        self.nodata = np.zeros(size, dtype=np.float64)
        self.Wasser = np.zeros(size, dtype=np.float64)
        self.Schill = np.zeros(size, dtype=np.float64)
        self.Muschel = np.zeros(size, dtype=np.float64)
        self.dense2 = np.zeros(size, dtype=np.float64)
        self.dense1 = np.zeros(size, dtype=np.float64)
        self.Strand = np.zeros(size, dtype=np.float64)
        self.Sand = np.zeros(size, dtype=np.float64)
        self.Misch = np.zeros(size, dtype=np.float64)
        self.Misch2 = np.zeros(size, dtype=np.float64)
        self.Schlick = np.zeros(size, dtype=np.float64)
        self.schlick_t = np.zeros(size, dtype=np.float64)
        self.Wasser2 = np.zeros(size, dtype=np.float64)

@jit(nopython=True)
def apply_rules(inputs: Inputs, outputs: Outputs):
    for i in range(len(outputs.nodata)):
        t0 = 1.0
        #    if b4 is NODATA:
        t1 = min(t0, _B4_NODATA(inputs.b4[i]))
        #        nodata = TRUE
        outputs.nodata[i] = t1
        #    elif (b8 is GT_033 and b1 is LT_085) or b8 is LT_009:
        t1 = min(t0, 1.0 - t1)
        t2 = min(t1, max(min(_B8_GT_033(inputs.b8[i]), _B1_LT_085(inputs.b1[i])), _B8_LT_009(inputs.b8[i])))
        #        if b5 is LT_01:
        t3 = min(t2, _B5_LT_01(inputs.b5[i]))
        #            Wasser = TRUE
        outputs.Wasser[i] = t3
        #        elif (b19 is GT_015 and (b8 is GT_04 and b8 is LT_085) and b7 is LT_05) or (b8 is GT_04 and bsum is GT_011) or (b8 is GT_035 and bsum is GT_016):
        t3 = min(t2, 1.0 - t3)
        t4 = min(t3, max(max(min(min(_B19_GT_015(inputs.b19[i]), min(_B8_GT_04(inputs.b8[i]), _B8_LT_085(inputs.b8[i]))), _B7_LT_05(inputs.b7[i])), min(_B8_GT_04(inputs.b8[i]), _BSum_GT_011(inputs.bsum[i]))), min(_B8_GT_035(inputs.b8[i]), _BSum_GT_016(inputs.bsum[i]))))
        #            if bsum is GT_013:
        t5 = min(t4, _BSum_GT_013(inputs.bsum[i]))
        #                Schill = TRUE
        outputs.Schill[i] = t5
        #            else:
        t5 = min(t4, 1.0 - t5)
        #                Muschel = TRUE
        outputs.Muschel[i] = t5
        #        elif b8 is GT_045:
        t4 = min(t3, 1.0 - t4)
        t5 = min(t4, _B8_GT_045(inputs.b8[i]))
        #            dense2 = TRUE
        outputs.dense2[i] = t5
        #        else:
        t5 = min(t4, 1.0 - t5)
        #            dense1 = TRUE
        outputs.dense1[i] = t5
        #    elif b1 is GT_1:
        t2 = min(t1, 1.0 - t2)
        t3 = min(t2, _B1_GT_1(inputs.b1[i]))
        #        Strand = TRUE
        outputs.Strand[i] = t3
        #    elif b3 is LT_005:
        t3 = min(t2, 1.0 - t3)
        t4 = min(t3, _B3_LT_005(inputs.b3[i]))
        #        Sand = TRUE
        outputs.Sand[i] = t4
        #    elif b3 is LT_01 and b8 is GT_0:
        t4 = min(t3, 1.0 - t4)
        t5 = min(t4, min(_B3_LT_01(inputs.b3[i]), _B8_GT_0(inputs.b8[i])))
        #        Misch = TRUE
        outputs.Misch[i] = t5
        #    elif b3 is LT_015 and b8 is GT_0:
        t5 = min(t4, 1.0 - t5)
        t6 = min(t5, min(_B3_LT_015(inputs.b3[i]), _B8_GT_0(inputs.b8[i])))
        #        Misch2 = TRUE
        outputs.Misch2[i] = t6
        #    elif b3 is LT_02 and b2 is GT_0 and b8 is GT_0:
        t6 = min(t5, 1.0 - t6)
        t7 = min(t6, min(min(_B3_LT_02(inputs.b3[i]), _B2_GT_0(inputs.b2[i])), _B8_GT_0(inputs.b8[i])))
        #        Schlick = TRUE
        outputs.Schlick[i] = t7
        #    elif b16 is GT_0 and b8 is GT_0:
        t7 = min(t6, 1.0 - t7)
        t8 = min(t7, min(_B16_GT_0(inputs.b16[i]), _B8_GT_0(inputs.b8[i])))
        #        schlick_t = TRUE
        outputs.schlick_t[i] = t8
        #    else:
        t8 = min(t7, 1.0 - t8)
        #        Wasser2 = TRUE
        outputs.Wasser2[i] = t8

max_stars_repo_name: forman/dectree

max_issues_repo_name: forman/dectree

max_forks_repo_name: forman/dectree

reason: import numpy,from numba