main.py

# -*- coding: utf-8 -*-
import os
import sys
import argparse
from glob import glob
import pandas as pd
from evaluation.metrics import get_metrics
from HNAD import runner

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
parser = argparse.ArgumentParser()
parser.add_argument('--data_root', type=str, default=f'./datasets/TSB-AD-U/')
parser.add_argument("--runs", type=int, default=1,
                    help="how many times we repeat the experiments to "
                         "obtain the average performance")
parser.add_argument("--dataset", type=str,
                    default='Stock',
                    help='dataset name or a list of names split by comma')
parser.add_argument('--epochs', type=int, default=50)
parser.add_argument('--batch_size', type=int, default=128)
parser.add_argument('--lr', type=float, default=8e-4)
parser.add_argument('--window_size', type=int, default=50)
parser.add_argument('--random_state', type=int, default=2024)
parser.add_argument('--lambda_spectral', type=float, default=1)
parser.add_argument('--n_classes', type=int, default=4)
parser.add_argument('--temperature', type=float, default=10)

args = parser.parse_args()
model_configs = {
    'epochs': args.epochs,
    'batch_size': args.batch_size,
    'lr': args.lr,
    'window_size': args.window_size,
    'lambda_spectral': args.n_classes,
    'n_classes': args.n_classes,
    'temperature': args.temperature
}
dataset_name_lst = args.dataset.split(',')

for dataset in dataset_name_lst:
    for path in glob(f'{args.data_root}/*{dataset}*.csv'):
        print('Processing:{} by HNAD'.format(dataset))

        df = pd.read_csv(path).dropna()
        data = df.iloc[:, 0:-1].values.astype(float)
        label = df['Label'].astype(int).to_numpy()

        train_index = path.split('/')[-1].split('.')[0].split('_')[-3]
        data_train = data[:int(train_index), :]
        output = runner(data_train, data, seed=args.random_state, **model_configs)

        evaluation_result = get_metrics(output, label, slidingWindow=args.window_size)
        print('Evaluation Result: ', evaluation_result)