ablation.py

import os
import logging
import hydra
from omegaconf import DictConfig, OmegaConf
from copy import deepcopy
import numpy as np

from scipy.stats import binomtest, chi2

from train import run_experiment

log = logging.getLogger(__name__)

def save_txt(filename,content):
    os.makedirs(os.path.dirname(filename),exist_ok=True)
    with open(filename,'a') as f:
        f.write(content+'\n')

def _mcnemar_pvalue(y_true,y_pred_A,y_pred_B,exact_threshold=25):
    y_true = np.asarray(y_true)
    yA = np.asarray(y_pred_A)
    yB = np.asarray(y_pred_B)

    A_correct = (yA == y_true)
    B_correct = (yB == y_true)

    b= int(np.sum(A_correct & (~B_correct)))
    c= int(np.sum((~A_correct) & B_correct))
    n_disagree = b+c

    if n_disagree == 0:
        return 1.0,b,c
    
    if n_disagree < exact_threshold:
        p=binomtest(b,n_disagree,0.5,alternative='two-sided').pvalue
    else:
        chi2_stat = (abs(b - c) - 1)**2 / n_disagree
        p = chi2.sf(chi2_stat, df=1)
    
    return float(p),b,c

def ablation_jc(base_cfg):
    log.info("=== Ablation: Jc ===")

    jcs=[0.005, 0.01, 0.015, 0.02, 0.03, 0.05, 0.1, 0.2, 0.3, 0.5, 0.7, 0.9]
    output_file = f"results/ablation/jc_{base_cfg.data.dataset}_{base_cfg.attack.mode}_{base_cfg.attack.rate}.txt"
    save_txt(output_file,'jc_list:')

    for jc in jcs:
        cfg = deepcopy(base_cfg)
        cfg.model.threshold = jc
        
        try:
            res = run_experiment(cfg)
            line = f"jc: {jc}, acc: {res['acc']}, precision: {res['pre']}, f1: {res['f1']}"
            log.info(line)
            save_txt(output_file, line)
        except Exception as e:
            log.error(f"Error at jc={jc}: {e}")
    save_txt(output_file, '--------------------------')

def ablation_jl(base_cfg):
    log.info("=== Ablation: Jl (Link Threshold) ===")
    
    jls = [0.005, 0.01, 0.015, 0.02, 0.03, 0.05, 0.1, 0.2, 0.3, 0.5, 0.7, 0.9]
    output_file = f'results/ablation/jl_{base_cfg.data.dataset}_{base_cfg.attack.mode}_{base_cfg.attack.rate}.txt'
    save_txt(output_file, 'jl list:')
    
    for jl in jls:
        cfg = deepcopy(base_cfg)
        cfg.model.threshold2 = jl
        
        try:
            res = run_experiment(cfg)
            line = f"jl: {jl}, acc: {res['acc']}, precision: {res['pre']}, f1: {res['f1']}"
            log.info(line)
            save_txt(output_file, line)
        except Exception as e:
            log.error(f"Error at jl={jl}: {e}")
    save_txt(output_file, '--------------------------')

def ablation_time(base_cfg):
    log.info("=== Ablation: Time Comparison ===")
    datasets = ['coauthorship_cora', 'cocitation_citeseer', 'cocitation_cora', 'dblp4k-paper']
    models = ['gcnj', 'rgcn', 'gcnsvd', 'elasso', 'hgnn_shield']
    
    output_file = 'results/ablation/time_results.txt'
    save_txt(output_file, 'time comparison')
    
    for data in datasets:
        save_txt(output_file, f'data {data}')
        for m in models:
            cfg = deepcopy(base_cfg)
            cfg.data.dataset = data
            cfg.model.name = m
            
            try:
                res = run_experiment(cfg)
                line = f"model {m}, train {res['train_time']}, infer {res['infer_time']}, acc {res['acc']}, f1 {res['f1']}"
                log.info(line)
                save_txt(output_file, line)
            except Exception as e:
                log.error(f"Error {m} on {data}: {e}")
        save_txt(output_file, '---------')

def ablation_pvalue(base_cfg):
    log.info("=== Ablation: P-Value Analysis ===")
    models = ['elasso', 'hgnn_shield']
    datasets = ['coauthorship_cora', 'cocitation_citeseer', 'cocitation_cora', 'dblp4k-paper']
    
    output_file = f'results/ablation/pvalue_results.txt'
    header = "dataset\tmodel_A\tacc_A\tmodel_B\tacc_B\tb\tc\tpvalue"
    save_txt(output_file, header)
    
    for data in datasets:
        results = {}
        for m in models:
            cfg = deepcopy(base_cfg)
            cfg.data.dataset = data
            cfg.model.name = m
            try:
                log.info(f"Running {m} on {data}...")
                results[m] = run_experiment(cfg)
            except Exception as e:
                log.error(f"Error {m} on {data}: {e}")
                results[m] = None
        
        res_A = results[models[0]]
        res_B = results[models[1]]
        
        if res_A and res_B:
            y_true = res_A['y_true']
            y_pred_A = res_A['y_pred']
            y_pred_B = res_B['y_pred']
            
            pval, b, c = _mcnemar_pvalue(y_true, y_pred_A, y_pred_B)
            
            line = f"{data}\t{models[0]}\t{res_A['acc']:.4f}\t{models[1]}\t{res_B['acc']:.4f}\t{b}\t{c}\t{pval:.6e}"
            log.info(f"Result: {line}")
            save_txt(output_file, line)

@hydra.main(config_path="config", config_name="config", version_base="1.2")
def main(cfg: DictConfig):
    task = cfg.ablation.task
    
    if task == 'jc':
        ablation_jc(cfg)
    elif task == 'jl':
        ablation_jl(cfg)
    elif task == 'time':
        ablation_time(cfg)
    elif task == 'pvalue':
        ablation_pvalue(cfg)
    else:
        log.error(f"Unknown ablation task: {task}")

if __name__ == '__main__':
    main()