loop.py

import os
import time
import pandas as pd
from noise import Noise
from explain import Explain
from model import Model
from autogluon.tabular import TabularDataset

class Loop:
    def __init__(self, model_name='RandomForestGini', noise_type='gaussian', standard=0.05, importance_type='default', min_bound="0.95"):
        self.noise = Noise(noise_type, standard)
        self.explainer = Explain(importance_type)
        self.model = Model(model_name)

        self.origin_train = None
        self.origin_val = None
        self.origin_test = None

        self.importance_type = importance_type
        self.feature_noise = {'base': 0}
        self.feature_importance_dict = {'base': 0.0}

        
        self.forbidden_feature = set()
        self.best_score = 0
        self.best_ece = 0
        self.min_bound = min_bound
        self.iteration_times = []

    def set_data(self, train_path, val_path=None, test_path=None):
        if train_path: self.origin_train = pd.read_csv(train_path)
        if val_path: self.origin_val = pd.read_csv(val_path)
        if test_path: self.origin_test = pd.read_csv(test_path)

    def iterate_window(self, model_path, results_path, iid_data, visible_ood, invisible_ood, num_gpus, rounds=10, window=5, mode="soft"):
        threshold = self.min_bound if "soft" in mode else 1.0
        base_window = window
        for i in range(rounds):
            print(f"\n===== Iteration {i+1}/{rounds} =====")
            select_feature_elapsed = 0.0
            base_window_elapsed = 0.0
            # saving_path = os.path.join(model_path, 'base_model' if i == 0 else f'iter_{i+1}')
            saving_path = model_path
            if i == 0:
                base_round_start = time.perf_counter()
                print(f"[Noise Update] {self.feature_noise}")
                noised_train, noised_val = self.noise.apply(self.origin_train, self.origin_val, self.feature_noise)
                
                best_predictor = self.model.run(noised_train, noised_val, saving_path, num_gpus)
                best_visible_results_df = self.model.evaluate(visible_ood, "Visible OOD")
                ood_acc = best_visible_results_df['score_test'].values[0]

                print(f"[BASE] OOD accuracy {ood_acc:.4f} > {threshold:.2f} × best {self.best_score:.4f}")
                
                base_round_elapsed = time.perf_counter() - base_round_start
                print(f"[Timing] Base round took {base_round_elapsed:.4f}s")
            else:
                
                top_features = sorted(self.feature_importance_dict.items(), key=lambda x: x[1], reverse=True)
                key_feature = None
                best_visible_results_df = None
                least_decrease = -100.0

                select_feature_start = time.perf_counter()
                for index, (f, _ )in enumerate(top_features):
                    if index < window:
                        print(f"             ===== Iteration {i+1}-{index+1}/{f} =====")
                        self.increase_feature_noise(f)

                        noised_train, noised_val = self.noise.apply(self.origin_train, self.origin_val, self.feature_noise)
                        predictor = self.model.run(noised_train, noised_val, saving_path, num_gpus)
                        visible_df = self.model.evaluate(visible_ood, "Visible OOD")
                        ood_acc = visible_df['score_test'].values[0]

                        if ood_acc > self.best_score * threshold:
                            decrease = ood_acc - self.best_score
                            print(f"[Window] Feature {f}, OOD accuracy: {ood_acc:.4f}, Best accuracy: {self.best_score:.4f}, Decrease: {decrease:.4f}")
                            if decrease > least_decrease:
                                least_decrease = decrease
                                key_feature = f
                                print(f"[Window Best] Best Feature {key_feature} least decrease updated to {least_decrease:.4f}")
                        else:
                            print(f"[Window] Feature {f} did not improve OOD accuracy: {ood_acc:.4f} ≤ best {self.best_score:.4f}")
                        
                        self.decrease_feature_noise(f)

                        if index + 1 == window and key_feature is None:
                            print("[Window] No feature improved OOD accuracy in this window. Extend Window.")
                            window += 1
                        elif index + 1 == base_window:
                            print(f"[Noise Update] Selected feature {key_feature} with least decrease {least_decrease:.4f}")
                            self.increase_feature_noise(key_feature)

                            select_feature_elapsed = time.perf_counter() - select_feature_start
                            print(f"[Timing] Select feature took {select_feature_elapsed:.4f}s")

                            # need to re-run best model to cover, otherwise, there
                            #is a conflict between self.feature and best_predictor
                            
                            base_window_start = time.perf_counter()
                            print(f"[Re-evaluate] Re-evaluating best predictor after selecting feature {key_feature}")
                            noised_train, noised_val = self.noise.apply(self.origin_train, self.origin_val, self.feature_noise)
                            best_predictor = self.model.run(noised_train, noised_val, saving_path, num_gpus)
                            best_visible_results_df = self.model.evaluate(visible_ood, "Visible OOD")

                            window = base_window
                            base_window_elapsed = time.perf_counter() - base_window_start
                            print(f"[Timing] Base window re-run took {base_window_elapsed:.4f}s")
                            break
                                


            # 不知道这个会对模型训练有什么影响，之前一直没更新best score
            get_importance_start = time.perf_counter()
            self.best_score = best_visible_results_df['score_test'].values[0]
            # print(best_predictor._learner.feature_generator.feature_metadata_in.get_features())
            self.feature_importance_dict = self.explainer.get_importance(best_predictor, self.origin_val)
            print(f"\n[Feature Importance] Top 5: {sorted(self.feature_importance_dict.items(), key=lambda x: x[1], reverse=True)[:5]}")
            get_importance_elapsed = time.perf_counter() - get_importance_start
            print(f"[Timing] Get importance took {get_importance_elapsed:.4f}s")

            evaluate_verification_start = time.perf_counter()
            iid_df = self.model.evaluate(iid_data, name="IID")
            self.save_results(iid_df, f'{results_path}/leaderboard-False-{self.noise.standard}.csv')
            self.save_results(best_visible_results_df, f'{results_path}/leaderboard-True-visible-{self.noise.standard}.csv')
            self.save_importances(results_path)
            evaluate_verification_elapsed = time.perf_counter() - evaluate_verification_start
            print(f"[Timing] Evaluate verification took {evaluate_verification_elapsed:.4f}s")


            invisible_df = self.model.evaluate(invisible_ood, name="Invisible OOD")
            self.save_results(invisible_df, f'{results_path}/leaderboard-True-invisible-{self.noise.standard}.csv')

            # 这里并没有考虑到所有特征都被禁止的情况，当window大于所含有的重要特征，即认为收敛
            # print(f"[Window Adjust] Reducing base window to {base_window} and current window to {window}.")
            # print(len(self.feature_importance_dict))

            iteration_time_record = {
                "iteration": i + 1,
                "base_time": base_round_elapsed if i == 0 else base_window_elapsed,
                "selection_time": 0.0 if i == 0 else select_feature_elapsed,
                "explain_time": get_importance_elapsed,
                "verification_time": evaluate_verification_elapsed,
            }
            self.save_iteration_time(results_path, iteration_time_record)

            if window == 2:
                print("[Forbidden] All features forbidden or evaluated. Stopping.")
                break
            elif window > len(self.feature_importance_dict):
                base_window = base_window - 1
                window = base_window
                print(f"[Window Adjust] Reducing base window to {base_window} and current window to {window}.")

        print("[Finished] Experiment complete.")

    def save_iteration_time(self, results_path, iteration_time_record):
        self.iteration_times.append(iteration_time_record)

        filename = f'{results_path}/iteration_times.csv'
        df = pd.DataFrame([iteration_time_record])
        header = not os.path.exists(filename)
        df.to_csv(filename, mode='a', header=header, index=False)

        print(
            "[Timing Summary] "
            f"iter={iteration_time_record['iteration']} "
            f"base_time={iteration_time_record['base_time']:.4f}s "
            f"selection_time={iteration_time_record['selection_time']:.4f}s "
            f"explain_time={iteration_time_record['explain_time']:.4f}s "
            f"verification_time={iteration_time_record['verification_time']:.4f}s"
        )

    def save_results(self, df, path):
        df.insert(1, 'feature', '+'.join(self.feature_noise.keys()))
        df.insert(2, 'degree', '+'.join(map(str, self.feature_noise.values())))
        header = not os.path.exists(path)
        df.to_csv(path, index=False, mode='a', header=header)
        print(f"[Save] Results saved to {path}")

    def save_importances(self, path):
        filename = f'{path}/feature_importance.csv'
        
        df = pd.DataFrame([self.feature_importance_dict])
        df.insert(0, 'feature', '+'.join(self.feature_noise.keys()))
        df.insert(1, 'degree', '+'.join(map(str, self.feature_noise.values())))

        if os.path.exists(filename):
            existing = pd.read_csv(filename, nrows=0)
            existing_cols = existing.columns.tolist()
            new_cols = df.columns.tolist()

            all_cols = existing_cols.copy()
            for c in new_cols:
                if c not in all_cols:
                    all_cols.append(c)
            for c in all_cols:
                if c not in df.columns:
                    df[c] = 0
            
            df = df[all_cols]
            df.to_csv(filename, mode='a', header=False, index=False)
        else:
            df.to_csv(filename, mode='w', header=True, index=False)
        print(f"[Save] Feature importances saved to {filename}")


    def increase_feature_noise(self, feature):
        """Increase noise for feature."""
        self.feature_noise[feature] = self.feature_noise.get(feature, 0) + 1

    def decrease_feature_noise(self, feature):
        """Increase noise for feature."""
        if self.feature_noise[feature] == 1:
            del self.feature_noise[feature]
        else:
            self.feature_noise[feature] -= 1

    def update_forbidden_feature(self, feature):
        """Reduce noise for feature and add it to forbidden list."""
        self.decrease_feature_noise(feature)
        self.forbidden_feature.add(feature)