aiops/ContraLSP/attribution/mask_group.py

import time

import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
import torch
import torch.optim as optim
import pandas as pd
from attribution.mask import Mask
from attribution.perturbation import Perturbation
from tqdm import tqdm


class MaskGroup:

    def __init__(
        self,
        perturbation: Perturbation,
        device,
        random_seed: int = 987,
        deletion_mode: bool = False,
        verbose: bool = True,
    ):
        self.perturbation = perturbation
        self.device = device
        self.random_seed = random_seed
        self.verbose = verbose
        self.deletion_mode = deletion_mode
        self.mask_list = None
        self.area_list = None
        self.f = None
        self.X = None
        self.n_epoch = None
        self.T = None
        self.N_features = None
        self.Y_target = None
        self.masks_tensor = None
        self.mask_tensor = None
        self.hist = None

    def fit(
        self,
        X,
        f,
        area_list,
        loss_function,
        n_epoch: int = 1000,
        initial_mask_coeff: float = 0.5,
        size_reg_factor_init: float = 0.1,
        size_reg_factor_dilation: float = 100,
        learning_rate: float = 0.1,
        momentum: float = 0.9,
        time_reg_factor: float = 0,
    ):

        # Ensure that the area list is sorted
        area_list.sort()
        self.area_list = area_list
        N_area = len(area_list)
        # Create a list of masks
        mask_list = []
        # Initialize the random seed and the attributes
        t_fit = time.time()
        torch.manual_seed(self.random_seed)
        reg_factor = size_reg_factor_init
        error_factor = 1 - 2 * self.deletion_mode  # In deletion mode, the error has to be maximized
        reg_multiplicator = np.exp(np.log(size_reg_factor_dilation) / n_epoch)
        self.f = f
        self.X = X
        self.n_epoch = n_epoch
        self.T, self.N_features = X.shape
        self.Y_target = f(X)
        # The initial mask tensor has all coefficients set to initial_mask_coeff
        self.masks_tensor = initial_mask_coeff * torch.ones(size=(N_area, self.T, self.N_features), device=self.device)
        # The target is the same for each mask so we simply repeat it along the first axis
        Y_target_group = self.Y_target.clone().detach().unsqueeze(0).repeat(N_area, 1, 1)
        # Create a copy of the extremal tensor that is going to be trained, the optimizer and the history
        masks_tensor_new = self.masks_tensor.clone().detach().requires_grad_(True)
        optimizer = optim.SGD([masks_tensor_new], lr=learning_rate, momentum=momentum)
        hist = torch.zeros(3, 0)
        # Initializing the reference vector used in the regulator
        reg_ref = torch.ones((N_area, self.T * self.N_features), dtype=torch.float32, device=self.device)
        for i, area in enumerate(self.area_list):
            reg_ref[i, : int((1 - area) * self.T * self.N_features)] = 0.0
        # Run the optimization
        for k in range(n_epoch):
            # Measure the loop starting time
            t_loop = time.time()
            # Generate perturbed input and outputs
            if self.deletion_mode:
                X_pert = self.perturbation.apply_extremal(X=X, extremal_tensor=1 - masks_tensor_new)
            else:
                X_pert = self.perturbation.apply_extremal(X=X, extremal_tensor=masks_tensor_new)
            Y_pert = torch.stack([f(x_pert) for x_pert in torch.unbind(X_pert, dim=0)], dim=0)

            # Evaluate the overall loss (error [L_e] + size regulation [L_a] + time variation regulation [L_c])
            error = loss_function(Y_pert, Y_target_group)
            masks_tensor_sorted = masks_tensor_new.reshape(N_area, self.T * self.N_features).sort(dim=1)[0]
            size_reg = ((reg_ref - masks_tensor_sorted) ** 2).mean()
            time_reg = (torch.abs(masks_tensor_new[:, 1 : self.T - 1, :] - masks_tensor_new[:, : self.T - 2, :])).mean()
            loss = error_factor * error + reg_factor * size_reg + time_reg_factor * time_reg
            # Apply the gradient step
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            # Ensures that the constraint is fulfilled
            masks_tensor_new.data = masks_tensor_new.data.clamp(0, 1)
            # Save the error and the regulator
            metrics = torch.tensor([error, size_reg, time_reg]).cpu().unsqueeze(1)
            hist = torch.cat((hist, metrics), dim=1)
            # Increase the regulator coefficient
            reg_factor *= reg_multiplicator
            # Measure the loop ending time
            t_loop = time.time() - t_loop
            if self.verbose and k%200==0:
                print(
                    f"Epoch {k + 1}/{n_epoch}: error = {error.data:.3g} ; "
                    f"size regulator = {size_reg.data:.3g} ; time regulator = {time_reg.data:.3g} ;"
                    f" time elapsed = {t_loop:.3g} s"
                )
                print((masks_tensor_new.sum()).data/N_area)

        # Update the mask and history tensor, print the final message
        self.masks_tensor = masks_tensor_new.clone().detach().requires_grad_(False)
        self.hist = hist
        t_fit = time.time() - t_fit
        print(
            f"The optimization finished: error = {error.data:.3g} ; size regulator = {size_reg.data:.3g} ;"
            f" time regulator = {time_reg.data:.3g} ; time elapsed = {t_fit:.3g} s"
        )

        # Store the individual mask coefficients in distinct mask objects
        for index, mask_tensor in enumerate(self.masks_tensor.unbind(dim=0)):
            mask = Mask(
                perturbation=self.perturbation, device=self.device, verbose=False, deletion_mode=self.deletion_mode
            )
            mask.mask_tensor = mask_tensor
            mask.hist = self.hist
            mask.f = self.f
            mask.X = self.X
            mask.n_epoch = self.n_epoch
            mask.T, mask.N_features = self.T, self.N_features
            mask.Y_target = self.Y_target
            mask.loss_function = loss_function
            mask_list.append(mask)
        self.mask_list = mask_list

    def fit_multiple(
        self,
        X,
        f,
        area_list,
        loss_function_multiple,
        use_last_timestep_only: bool = False,
        n_epoch: int = 200,
        initial_mask_coeff: float = 0.5,
        size_reg_factor_init: float = 0.1,
        size_reg_factor_dilation: float = 100,
        learning_rate: float = 0.1,
        momentum: float = 0.9,
        time_reg_factor: float = 0.01,
    ):
        # Ensure that the area list is sorted
        area_list.sort()
        self.area_list = area_list
        N_area = len(area_list)
        # Create a list of masks
        mask_list = []
        # Initialize the random seed and the attributes
        t_fit = time.time()
        torch.manual_seed(self.random_seed)
        reg_factor = size_reg_factor_init
        error_factor = 1 - 2 * self.deletion_mode  # In deletion mode, the error has to be maximized
        reg_multiplicator = np.exp(np.log(size_reg_factor_dilation) / n_epoch)
        self.f = f
        self.X = X
        self.n_epoch = n_epoch
        num_samples, self.T, self.N_features = X.shape
        self.Y_target = f(X) # num_samples, num_time, num_state=2
        if use_last_timestep_only:
            self.Y_target = self.Y_target[:, -1:, :]
        # The initial mask tensor has all coefficients set to initial_mask_coeff
        self.masks_tensor = initial_mask_coeff * torch.ones(size=(N_area, num_samples, self.T, self.N_features), device=self.device)
        # The target is the same for each mask so we simply repeat it along the first axis
        Y_target_group = self.Y_target.clone().detach().unsqueeze(0).repeat(N_area, 1, 1, 1)
        # Create a copy of the extremal tensor that is going to be trained, the optimizer and the history
        masks_tensor_new = self.masks_tensor.clone().detach().requires_grad_(True)
        # optimizer = optim.SGD([masks_tensor_new], lr=learning_rate, momentum=momentum)
        optimizer = optim.Adam([masks_tensor_new], lr=learning_rate)

        metrics = []
        # Initializing the reference vector used in the regulator
        reg_ref = torch.ones((N_area, num_samples, self.T * self.N_features), dtype=torch.float32, device=self.device)
        for i, area in enumerate(self.area_list):
            reg_ref[i, :, :int((1 - area) * self.T * self.N_features)] = 0.0

        # Run the optimization
        for k in tqdm(range(n_epoch)):
            # Measure the loop starting time
            t_loop = time.time()
            # Generate perturbed input and outputs
            if self.deletion_mode:
                X_pert = self.perturbation.apply_extremal_multiple(X=X, extremal_tensor=1 - masks_tensor_new)
            else:
                X_pert = self.perturbation.apply_extremal_multiple(X=X, extremal_tensor=masks_tensor_new)

            # x_pert (num_samples, T, num_feature)
            # f(x_pert) = (num_sample, T, num_state)
            # y_pert = (num_area, num_sample, T, num_state)

            # x_pert (T, num_feature)
            # f(x_pert) = (T, num_state)
            # Y_pert = (n_area, T, num_state)
            X_pert_flatten = X_pert.reshape(N_area * num_samples, self.T, self.N_features)
            Y_pert_flatten = f(X_pert_flatten) # (N_area * num_samples, T, num_state)
            if use_last_timestep_only:
                Y_pert = Y_pert_flatten.reshape(N_area, num_samples, 1, -1)
            else:
                Y_pert = Y_pert_flatten.reshape(N_area, num_samples, self.T, -1)

            # Evaluate the overall loss (error [L_e] + size regulation [L_a] + time variation regulation [L_c])
            error = loss_function_multiple(Y_pert, Y_target_group) # (num_sample)
            masks_tensor_sorted = masks_tensor_new.reshape(N_area, num_samples, self.T * self.N_features).sort(dim=2)[0]
            size_reg = ((reg_ref - masks_tensor_sorted) ** 2).mean(dim=[0, 2])
            masks_tensor_diff = masks_tensor_new[:, :, 1: self.T - 1, :] - masks_tensor_new[:, :, :self.T - 2, :]
            time_reg = (torch.abs(masks_tensor_diff)).mean(dim=[0, 2, 3])
            loss = error_factor * error + reg_factor * size_reg + time_reg_factor * time_reg
            # Apply the gradient step
            optimizer.zero_grad()
            loss.backward(torch.ones_like(loss))
            optimizer.step()
            # Ensures that the constraint is fulfilled
            masks_tensor_new.data = masks_tensor_new.data.clamp(0, 1)
            # Save the error and the regulator
            metric = torch.stack([error, size_reg, time_reg], dim=1).detach().cpu().numpy()
            metrics.append(metric)
            # Increase the regulator coefficient
            reg_factor *= reg_multiplicator
            # Measure the loop ending time
            t_loop = time.time() - t_loop
            if self.verbose and k%20==0:
                print(
                    f"Epoch {k + 1}/{n_epoch}: error = {error.mean().data:.3g} ; "
                    f"size regulator = {size_reg.mean().data:.3g} ; time regulator = {time_reg.mean().data:.3g} ;"
                    f" time elapsed = {t_loop:.3g} s"
                )
                print((masks_tensor_new.sum()/(N_area*num_samples)).data, X.shape[2]*X.shape[1])

        # Update the mask and history tensor, print the final message
        self.masks_tensor = masks_tensor_new.clone().detach().requires_grad_(False) # (N_area, num_samples, T, nfeat)
        self.hist = torch.from_numpy(np.stack(metrics, axis=2))
        t_fit = time.time() - t_fit

        print(
            f"The optimization finished: error = {error.mean().data:.3g} ; size regulator = {size_reg.mean().data:.3g} ;"
            f" time regulator = {time_reg.mean().data:.3g} ; time elapsed = {t_fit:.3g} s"
        )

        # Store the individual mask coefficients in distinct mask objects
        for index, mask_tensor in enumerate(self.masks_tensor.unbind(dim=0)):
            mask = Mask(
                perturbation=self.perturbation, device=self.device, verbose=False, deletion_mode=self.deletion_mode
            )
            mask.mask_tensor = mask_tensor
            mask.hist = self.hist
            mask.f = self.f
            mask.X = self.X
            mask.n_epoch = self.n_epoch
            mask.T, mask.N_features = self.T, self.N_features
            mask.Y_target = self.Y_target
            mask.loss_function = loss_function_multiple
            mask_list.append(mask)
        self.mask_list = mask_list

    def get_best_mask(self):
        """This method returns the mask with lowest error."""
        error_list = [mask.get_error() for mask in self.mask_list]
        best_index = error_list.index(min(error_list))
        print(
            f"The mask of area {self.area_list[best_index]:.2g} is"
            f" the best with error = {error_list[best_index]:.3g}."
        )
        return self.mask_list[best_index]

    def get_extremal_mask(self, threshold):
        """This method returns the extremal mask for the acceptable error threshold (called epsilon in the paper)."""
        error_list = [mask.get_error() for mask in self.mask_list]
        # If the minimal error is above the threshold, the best we can do is select the mask with lowest error
        if min(error_list) > threshold:
            return self.get_best_mask()
        else:
            for id_mask, error in enumerate(error_list):
                if error < threshold:
                    print(
                        f"The mask of area {self.area_list[id_mask]:.2g} is"
                        f" extremal with error = {error_list[id_mask]:.3g}."
                    )
                    return self.mask_list[id_mask]

    def get_extremal_mask_multiple(self, thresholds):
        """This method returns the extremal mask for the acceptable error threshold (called epsilon in the paper)."""
        error_list = torch.stack([mask.get_error_multiple() for mask in self.mask_list], dim=1)
        mask_stacked = torch.stack([mask.mask_tensor for mask in self.mask_list])
        num_area, num_samples, num_times, num_features = mask_stacked.shape
        # If the minimal error is above the threshold, the best we can do is select the mask with lowest error
        thres_mask = torch.min(error_list, dim=1)[0] > thresholds
        best_mask = torch.argmin(error_list, dim=1) #(num_sample)
        error_mask = (error_list < thresholds.view(-1, 1)) * torch.arange(-len(self.mask_list), 0).view(1, -1).to(self.device)
        first_mask = torch.argmin(error_mask, dim=1)
        indexes = torch.where(thres_mask, best_mask, first_mask) # (num_sample)
        selected_masks = torch.gather(mask_stacked, 0, indexes.view(1, num_samples, 1, 1).expand(1, num_samples, num_times, num_features))
        self.mask_tensor = selected_masks.reshape(num_samples, num_times, num_features)
        return selected_masks.reshape(num_samples, num_times, num_features) #(num_samples, num_times, num_features)

    def plot_errors(self):
        """This method plots the error as a function of the mask size."""
        sns.set()
        error_list = [mask.get_error() for mask in self.mask_list]
        plt.plot(self.area_list, error_list)
        plt.title("Errors for the various masks")
        plt.xlabel("Mask area")
        plt.ylabel("Error")
        plt.show()

    def get_smooth_mask(self, sigma=1):
        """This method smooths the mask tensor by applying a temporal Gaussian filter for each feature.

        Args:
            sigma: Width of the Gaussian filter.

        Returns:
            torch.Tensor: The smoothed mask.
        """
        # Define the Gaussian smoothing kernel
        T_axis = torch.arange(1, self.T + 1, dtype=int, device=self.device)
        T1_tensor = T_axis.unsqueeze(1).unsqueeze(2)
        T2_tensor = T_axis.unsqueeze(0).unsqueeze(2)
        kernel_tensor = torch.exp(-1.0 * (T1_tensor - T2_tensor) ** 2 / (2.0 * sigma ** 2))
        kernel_tensor = torch.divide(kernel_tensor, torch.sum(kernel_tensor, 0))
        kernel_tensor = kernel_tensor.repeat(1, 1, self.N_features)
        # Smooth the mask tensor by applying the kernel
        mask_tensor_smooth = torch.einsum("sti,si->ti", kernel_tensor, self.mask_tensor)
        return mask_tensor_smooth

    def extract_submask(self, mask_tensor, ids_time, ids_feature):
        """This method extracts a submask specified with specified indices.

        Args:
            mask_tensor: The tensor from which data should be extracted.
            ids_time: List of the times that should be extracted.
            ids_feature: List of the features that should be extracted.

        Returns:
            torch.Tensor: Submask extracted based on the indices.
        """
        # If no identifiers have been specified, we use the whole data
        if ids_time is None:
            ids_time = [k for k in range(self.T)]
        if ids_feature is None:
            ids_feature = [k for k in range(self.N_features)]
        # Extract the relevant data in the mask
        submask_tensor = mask_tensor.clone().detach().requires_grad_(False).cpu()
        submask_tensor = submask_tensor[ids_time, :]
        submask_tensor = submask_tensor[:, ids_feature]
        return submask_tensor

    def plot_mask(self, idx=0, ids_time=None, ids_feature=None, smooth: bool = False, sigma: float = 1.0):
        """This method plots (part of) the mask.

        Args:
            ids_time: List of the times that should appear on the plot.
            ids_feature: List of the features that should appear on the plot.
            smooth: True if the mask should be smoothed before plotting.
            sigma: Width of the smoothing Gaussian kernel.

        Returns:
            None
        """
        sns.set()
        # Smooth the mask if required
        if smooth:
            mask_tensor = self.get_smooth_mask(sigma)
        else:
            mask_tensor = self.mask_tensor
        # Extract submask from ids
        submask_tensor_np = self.extract_submask(mask_tensor[idx], ids_time, ids_feature).numpy()
        df = pd.DataFrame(data=np.transpose(submask_tensor_np), index=ids_feature, columns=ids_time)
        # Generate heatmap plot
        color_map = sns.diverging_palette(10, 133, as_cmap=True)
        heat_map = sns.heatmap(data=df, cmap=color_map, cbar_kws={"label": "Mask"}, vmin=0, vmax=1)
        plt.xlabel("Time")
        plt.ylabel("Feature Number")
        plt.title("Mask coefficients over time")
        plt.show()